Fix sgd (#2777)

* Only update momentum buffers for SGD if momentum is enabled

* update stock PT version

Author: zhuhaozhe

dependency_version.yml

@@ -28,14 +28,14 @@ oneCCL:
 protobuf:
   version: 3.20.3
 pytorch:
-  version: 2.4.0.dev20240401+cpu
+  version: 2.4.0.dev20240417+cpu
 torch-ccl:
   commit: ccl_torch_dev_0131
   repo: https://github.com/intel/torch-ccl.git
   version: 2.3.0+cpu
 torchaudio:
-  version: 2.2.0.dev20240401+cpu
+  version: 2.2.0.dev20240417+cpu
 torchvision:
-  version: 0.19.0.dev20240401+cpu
+  version: 0.19.0.dev20240417+cpu
 transformers:
   version: 4.38.1

intel_extension_for_pytorch/optim/_functional.py

@@ -526,10 +526,11 @@ def sgd_step(self, closure=None):
             fused=self.fused,
         )
 
-        # update momentum_buffers in state
-        for p, momentum_buffer in zip(params_with_grad, momentum_buffer_list):
-            state = self.state[p]
-            state["momentum_buffer"] = momentum_buffer
+        if group["momentum"] != 0:
+            # update momentum_buffers in state
+            for p, momentum_buffer in zip(params_with_grad, momentum_buffer_list):
+                state = self.state[p]
+                state["momentum_buffer"] = momentum_buffer
 
     return loss