Step for SGD with momentum with 'lr'
momentum_step(p, lr, grad_avg, ...)
p
lr
learning rate
grad_avg
grad average
...
additional arguments to pass
None