Currently the function is passed into .step() and called without arguments.
See this line of code 24.
However, you could pass arguments using lambda and most likely partial.
Here is a small example:
model = nn.Linear(20, 2)
loss_fn = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=1.)
input = torch.randn(1, 20)
target = torch.empty(1, dtype=torch.long).random_(2)
optimizer.step(lambda: closure(0, torch.randn(size=model.weight.size())))

I cannot comment of the correctness of the approach, since the paper is behind a pay wall and I’m currently out of office, so I cannot see it.