Merge pull request #199 from kozistr/refactor/format

[Refactor] Fix styles

pyproject.toml

@@ -60,7 +60,7 @@ black = [
     { version = "==23.3.0", python = ">=3.7,<3.8" },
     { version = "^23.7.0", python = ">=3.8"}
 ]
-ruff = "^0.0.278"
+ruff = "^0.0.284"
 pytest = "^7.4.0"
 pytest-cov = "^4.1.0"
 

pytorch_optimizer/optimizer/lomo.py

@@ -46,7 +46,7 @@ def __init__(
         self.grad_norms: List[torch.Tensor] = []
         self.clip_coef: Optional[float] = None
 
-        p0: torch.Tensor = list(self.model.parameters())[0]
+        p0: torch.Tensor = next(iter(self.model.parameters()))
 
         self.grad_func: Callable[[Any], Any] = (
             self.fuse_update_zero3() if hasattr(p0, 'ds_tensor') else self.fuse_update()

pytorch_optimizer/optimizer/rotograd.py

@@ -219,7 +219,7 @@ def to(self, *args, **kwargs):
         self.backbone.to(*args, **kwargs)
         for head in self.heads:
             head.to(*args, **kwargs)
-        return super(RotateOnly, self).to(*args, **kwargs)
+        return super().to(*args, **kwargs)
 
     def train(self, mode: bool = True) -> nn.Module:
         super().train(mode)
@@ -284,7 +284,7 @@ def backward(self, losses: Sequence[torch.Tensor], backbone_loss=None, **kwargs)
         if not self.training:
             raise AssertionError('Backward should only be called when training')
 
-        if self.iteration_counter == 0 or self.iteration_counter == self.burn_in_period:
+        if self.iteration_counter in (0, self.burn_in_period):
             for i, loss in enumerate(losses):
                 self.initial_losses[i] = loss.item()
 

requirements-dev.txt

@@ -19,12 +19,12 @@ networkx==3.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 numpy==1.21.1 ; python_full_version >= "3.7.2" and python_version < "3.8"
 numpy==1.24.4 ; python_version >= "3.8" and python_full_version < "4.0.0"
 packaging==23.1 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-pathspec==0.11.1 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-platformdirs==3.9.1 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+pathspec==0.11.2 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+platformdirs==3.10.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pluggy==1.2.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pytest-cov==4.1.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pytest==7.4.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-ruff==0.0.278 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+ruff==0.0.284 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 sympy==1.12 ; python_version >= "3.8" and python_full_version < "4.0.0"
 tomli==2.0.1 ; python_full_version >= "3.7.2" and python_version < "3.11"
 torch==1.13.1+cpu ; python_full_version >= "3.7.2" and python_version < "3.8"

tests/test_optimizers.py

@@ -476,6 +476,6 @@ def test_lomo_optimizer(precision, environment):
     if precision == 16:
         optimizer.clip_coef = 0.9
 
-    loss = sphere_loss(list(model.parameters())[0])
+    loss = sphere_loss(next(iter(model.parameters())))
     optimizer.grad_norm(loss)
     optimizer.fused_backward(loss, lr=0.1)