Fix momentum initialization with weight decay

apex/optim/sgd.py

@@ -108,7 +108,6 @@ class SGD(Optimizer):
                 if 'momentum_buffer' not in param_state:
                     first_run = True
                     buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
-                    buf.mul_(momentum).add_(p.grad.data)
                     momentums.append(buf)
                 else:
                     first_run = False

csrc/multi_tensor_sgd_kernel.cu

@@ -102,6 +102,9 @@ struct SGDFunctor
           if (momentum != 0.f) {
             if (!first_run) {
               incoming_moms[ii] = incoming_moms[ii] * momentum + (1.f - dampening) * incoming_grads[ii];
+            } else {
+              // initialize momentume to current incoming grads
+              incoming_moms[ii] = incoming_grads[ii];
             }
 
             if (nesterov) {