style: fix RUF021

kozistr · Nov 23, 2024 · f82fc9d · f82fc9d
1 parent e53d049
commit f82fc9d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_optimizer/optimizer/lomo.py b/pytorch_optimizer/optimizer/lomo.py
@@ -83,7 +83,7 @@ def func(x: Any) -> Any:
                 if not p.requires_grad or p.grad is None:
                     continue
 
-                if self.loss_scaler and self.loss_scaler.has_overflow_serial or has_overflow(p.grad):
+                if (self.loss_scaler and self.loss_scaler.has_overflow_serial) or has_overflow(p.grad):
                     p.grad = None
                     self.loss_scaler.has_overflow_serial = True
                     break
@@ -119,7 +119,7 @@ def func(x: torch.Tensor) -> torch.Tensor:
 
                 all_reduce(p.grad, op=ReduceOp.AVG, async_op=False)
 
-                if self.loss_scaler and self.loss_scaler.has_overflow_serial or has_overflow(p.grad):
+                if (self.loss_scaler and self.loss_scaler.has_overflow_serial) or has_overflow(p.grad):
                     p.grad = None
                     self.loss_scaler.has_overflow_serial = True
                     break