Hi,
That looks good to me.
One thing we usually do in core optimizer is to lazily create the buffers (like v_old). Namely only create it during the first step.
That way, you can do things like:
optimizer = SGD(model=model, learning_rate=0.01, momentum=0.9)
mode.cuda()
# Your training loop

With your current code, this would fail because the buffer will be on cpu and the weights on gpu.