fix: device mismatch

Author: kozistr

pytorch_optimizer/optimizer/adan.py

@@ -83,7 +83,7 @@ def get_global_gradient_norm(self) -> Union[torch.Tensor, float]:
         if self.defaults['max_grad_norm'] == 0.0:
             return 1.0
 
-        global_grad_norm = get_global_gradient_norm(self.param_groups, self.param_groups[0]['params'][0].device)
+        global_grad_norm = get_global_gradient_norm(self.param_groups)
         global_grad_norm.sqrt_().add_(self.defaults['eps'])
 
         return torch.clamp(self.defaults['max_grad_norm'] / global_grad_norm, max=1.0)

pytorch_optimizer/optimizer/alig.py

@@ -52,7 +52,7 @@ def reset(self):
     @torch.no_grad()
     def compute_step_size(self, loss: float) -> float:
         r"""Compute step_size."""
-        global_grad_norm = get_global_gradient_norm(self.param_groups, torch.device('cpu'))
+        global_grad_norm = get_global_gradient_norm(self.param_groups)
         global_grad_norm.add_(1e-6)
 
         return loss / global_grad_norm.item()

pytorch_optimizer/optimizer/dadapt.py

@@ -473,7 +473,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
         numerator_weighted = group['numerator_weighted']
 
         if group['step'] == 0:
-            group['g0_norm'] = get_global_gradient_norm(self.param_groups, device).sqrt_().item()
+            group['g0_norm'] = get_global_gradient_norm(self.param_groups).sqrt_().item()
         g0_norm = group['g0_norm']
 
         if g0_norm == 0:

pytorch_optimizer/optimizer/lamb.py

@@ -103,7 +103,7 @@ def get_global_gradient_norm(self) -> Union[torch.Tensor, float]:
         if self.defaults['max_grad_norm'] == 0.0:
             return 1.0
 
-        global_grad_norm = get_global_gradient_norm(self.param_groups, self.param_groups[0]['params'][0].device)
+        global_grad_norm = get_global_gradient_norm(self.param_groups)
         global_grad_norm.sqrt_().add_(self.defaults['eps'])
 
         return torch.clamp(self.defaults['max_grad_norm'] / global_grad_norm, max=1.0)

pytorch_optimizer/optimizer/utils.py

@@ -272,9 +272,9 @@ def l2_projection(parameters: PARAMETERS, max_norm: float = 1e2):
 
 
 @torch.no_grad()
-def get_global_gradient_norm(param_groups: List[Dict], device: torch.device) -> torch.Tensor:
+def get_global_gradient_norm(param_groups: List[Dict]) -> torch.Tensor:
     r"""Get global gradient norm."""
-    global_grad_norm = torch.zeros(1, dtype=torch.float32, device=device)
+    global_grad_norm = torch.zeros(1, dtype=torch.float32, device=param_groups[0]['params'][0].device)
 
     for group in param_groups:
         for p in group['params']: