Update

fegin · fegin · commit f63199fac0b4 · 2024-11-26T14:46:45.000-08:00
[ghstack-poisoned]
diff --git a/torch/distributed/checkpoint/state_dict.py b/torch/distributed/checkpoint/state_dict.py
@@ -611,7 +611,11 @@ def _init_optim_state(optim: torch.optim.Optimizer) -> None:
     for param_group in optim.param_groups:
         if "lr" in param_group:
             lrs.append(param_group["lr"])
-            param_group["lr"] = 0.0
+            param_group["lr"] = (
+                torch.tensor(0.0)
+                if isinstance(param_group["lr"], torch.Tensor)
+                else 0.0
+            )
     optim.step(closure=None)
     # Whether to recover the "lr" should not matter too much as we will
     # restore checkpointing later.
