done

jon-chuang · pytorchmergebot · commit 18d6b92ae0a3 · 2023-11-21T19:23:55.000Z
diff --git a/test/optim/test_optim.py b/test/optim/test_optim.py
@@ -795,10 +795,9 @@ def _test_derived_optimizers(self, optimizer_pairs_with_flags, flag, reduced_pre
             st_state = state[0]
             mt_state = state[1]
 
+            assert_eq_kwargs = {}
             if reduced_precision:
                 assert_eq_kwargs = {'atol': 1e-5, 'rtol': 1e-4}
-            else:
-                assert_eq_kwargs = {}
 
             for st_p, mt_p in zip(res[0], res[1]):
                 self.assertEqual(st_p, mt_p, **assert_eq_kwargs)
diff --git a/torch/optim/adagrad.py b/torch/optim/adagrad.py
@@ -50,7 +50,7 @@ def __init__(
         for group in self.param_groups:
             for p in group["params"]:
                 state = self.state[p]
-                state["step"] = torch.tensor(0.0, dtype=torch.float)
+                state["step"] = torch.tensor(0.0, dtype=torch.float32)
                 init_value = (
                     complex(initial_accumulator_value, initial_accumulator_value)
                     if torch.is_complex(p)
@@ -73,7 +73,7 @@ def __setstate__(self, state):
         )
         if not step_is_tensor:
             for s in state_values:
-                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float)
+                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float32)
 
     def share_memory(self):
         for group in self.param_groups:
diff --git a/torch/optim/adam.py b/torch/optim/adam.py
@@ -75,7 +75,7 @@ def __setstate__(self, state):
         step_is_tensor = (len(state_values) != 0) and torch.is_tensor(state_values[0]['step'])
         if not step_is_tensor:
             for s in state_values:
-                s['step'] = torch.tensor(float(s['step']), dtype=torch.float)
+                s['step'] = torch.tensor(float(s['step']), dtype=torch.float32)
 
     def _init_group(
         self,
@@ -103,9 +103,9 @@ def _init_group(
                     # Deliberately host `step` on CPU if both capturable and fused are off.
                     # This is because kernel launches are costly on CUDA and XLA.
                     state['step'] = (
-                        torch.zeros((), dtype=torch.float, device=p.device)
+                        torch.zeros((), dtype=torch.float32, device=p.device)
                         if group['capturable'] or group['fused']
-                        else torch.tensor(0.0, dtype=torch.float)
+                        else torch.tensor(0.0, dtype=torch.float32)
                     )
                     # Exponential moving average of gradient values
                     state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
diff --git a/torch/optim/adamax.py b/torch/optim/adamax.py
@@ -56,7 +56,7 @@ def __setstate__(self, state):
         )
         if not step_is_tensor:
             for s in state_values:
-                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float)
+                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float32)
 
     def _init_group(self, group, params_with_grad, grads, exp_avgs, exp_infs, state_steps):
         has_complex = False
@@ -73,7 +73,7 @@ def _init_group(self, group, params_with_grad, grads, exp_avgs, exp_infs, state_
 
             # State initialization
             if len(state) == 0:
-                state["step"] = torch.tensor(0.0, dtype=torch.float)
+                state["step"] = torch.tensor(0.0, dtype=torch.float32)
                 state["exp_avg"] = torch.zeros_like(
                     p, memory_format=torch.preserve_format
                 )
diff --git a/torch/optim/adamw.py b/torch/optim/adamw.py
@@ -85,7 +85,7 @@ def __setstate__(self, state):
         )
         if not step_is_tensor:
             for s in state_values:
-                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float)
+                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float32)
 
     def _init_group(
         self,
@@ -115,9 +115,9 @@ def _init_group(
                 # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
                 # This is because kernel launches are costly on CUDA and XLA.
                 state["step"] = (
-                    torch.zeros((), dtype=torch.float, device=p.device)
+                    torch.zeros((), dtype=torch.float32, device=p.device)
                     if group["capturable"] or group["fused"]
-                    else torch.tensor(0.0, dtype=torch.float)
+                    else torch.tensor(0.0, dtype=torch.float32)
                 )
                 # Exponential moving average of gradient values
                 state["exp_avg"] = torch.zeros_like(
diff --git a/torch/optim/asgd.py b/torch/optim/asgd.py
@@ -62,19 +62,19 @@ def __setstate__(self, state):
         )
         if not step_is_tensor:
             for s in state_values:
-                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float)
+                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float32)
         eta_is_tensor = (len(state_values) != 0) and torch.is_tensor(
             state_values[0]["eta"]
         )
         if not eta_is_tensor:
             for s in state_values:
-                s["eta"] = torch.tensor(s["eta"], dtype=torch.float)
+                s["eta"] = torch.tensor(s["eta"], dtype=torch.float32)
         mu_is_tensor = (len(state_values) != 0) and torch.is_tensor(
             state_values[0]["mu"]
         )
         if not mu_is_tensor:
             for s in state_values:
-                s["mu"] = torch.tensor(float(s["mu"]), dtype=torch.float)
+                s["mu"] = torch.tensor(float(s["mu"]), dtype=torch.float32)
 
     def _init_group(self, group, params_with_grad, grads, mus, axs, etas, state_steps):
         has_complex = False
@@ -89,9 +89,9 @@ def _init_group(self, group, params_with_grad, grads, mus, axs, etas, state_step
                 state = self.state[p]
                 # State initialization
                 if len(state) == 0:
-                    state["step"] = torch.zeros((), device=p.device, dtype=torch.float)
-                    state["eta"] = torch.tensor(group["lr"], device=p.device, dtype=torch.float)
-                    state["mu"] = torch.ones((), device=p.device, dtype=torch.float)
+                    state["step"] = torch.zeros((), device=p.device, dtype=torch.float32)
+                    state["eta"] = torch.tensor(group["lr"], device=p.device, dtype=torch.float32)
+                    state["mu"] = torch.ones((), device=p.device, dtype=torch.float32)
                     state["ax"] = torch.zeros_like(
                         p, memory_format=torch.preserve_format
                     )
diff --git a/torch/optim/nadam.py b/torch/optim/nadam.py
@@ -40,11 +40,11 @@ def __setstate__(self, state):
         step_is_tensor = (len(state_values) != 0) and torch.is_tensor(state_values[0]['step'])
         if not step_is_tensor:
             for s in state_values:
-                s['step'] = torch.tensor(float(s['step']), dtype=torch.float)
+                s['step'] = torch.tensor(float(s['step']), dtype=torch.float32)
         mu_product_is_tensor = (len(state_values) != 0) and torch.is_tensor(state_values[0]['mu_product'])
         if not mu_product_is_tensor:
             for s in state_values:
-                s['mu_product'] = torch.tensor(s['mu_product'], dtype=torch.float)
+                s['mu_product'] = torch.tensor(s['mu_product'], dtype=torch.float32)
 
     def _init_group(self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, mu_products, state_steps):
         has_complex = False
@@ -63,12 +63,12 @@ def _init_group(self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, mu_
                     # Deliberately host `step` and `mu_product` on CPU if capturable is False.
                     # This is because kernel launches are costly on CUDA and XLA.
                     state['step'] = (
-                        torch.zeros((), dtype=torch.float, device=p.device)
-                        if group['capturable'] else torch.tensor(0.0, dtype=torch.float)
+                        torch.zeros((), dtype=torch.float32, device=p.device)
+                        if group['capturable'] else torch.tensor(0.0, dtype=torch.float32)
                     )
                     state['mu_product'] = (
-                        torch.ones((), dtype=torch.float, device=p.device)
-                        if group['capturable'] else torch.tensor(1.0, dtype=torch.float)
+                        torch.ones((), dtype=torch.float32, device=p.device)
+                        if group['capturable'] else torch.tensor(1.0, dtype=torch.float32)
                     )
                     # Exponential moving average of gradient values
                     state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
diff --git a/torch/optim/radam.py b/torch/optim/radam.py
@@ -65,7 +65,7 @@ def __setstate__(self, state):
         )
         if not step_is_tensor:
             for s in state_values:
-                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float)
+                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float32)
 
     def _init_group(self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps):
         has_complex = False
@@ -80,7 +80,7 @@ def _init_group(self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, sta
                 state = self.state[p]
                 # Lazy state initialization
                 if len(state) == 0:
-                    state["step"] = torch.tensor(0.0, dtype=torch.float)
+                    state["step"] = torch.tensor(0.0, dtype=torch.float32)
                     # Exponential moving average of gradient values
                     state["exp_avg"] = torch.zeros_like(
                         p, memory_format=torch.preserve_format
diff --git a/torch/optim/sparse_adam.py b/torch/optim/sparse_adam.py
@@ -35,16 +35,6 @@ def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, maximize: bool
         defaults = dict(lr=lr, betas=betas, eps=eps, maximize=maximize)
         super().__init__(params, defaults)
 
-    def __setstate__(self, state):
-        super().__setstate__(state)
-        state_values = list(self.state.values())
-        step_is_tensor = (len(state_values) != 0) and torch.is_tensor(
-            state_values[0]["step"]
-        )
-        if not step_is_tensor:
-            for s in state_values:
-                s["step"] = torch.tensor(float(s["step"]), dtype=torch.float)
-
     @torch.no_grad()
     def step(self, closure=None):
         """Perform a single optimization step.
@@ -80,7 +70,7 @@ def step(self, closure=None):
 
                     # State initialization
                     if len(state) == 0:
-                        state['step'] = torch.tensor(0.0, dtype=torch.float)
+                        state['step'] = 0
                         # Exponential moving average of gradient values
                         state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                         # Exponential moving average of squared gradient values
