fix: type

kozistr · Dec 3, 2024 · 756d7ea · 756d7ea
1 parent 0274639
commit 756d7ea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_optimizer/optimizer/muon.py b/pytorch_optimizer/optimizer/muon.py
@@ -97,8 +97,8 @@ def __init__(
         adamw_params = self.get_parameters(adamw_params) if adamw_params is not None else []
         params.extend(adamw_params)
 
-        self.world_size: int = os.environ.get('WORLD_SIZE', 1)
-        self.rank: int = os.environ.get('RANK', 0)
+        self.world_size: int = int(os.environ.get('WORLD_SIZE', 1))
+        self.rank: int = int(os.environ.get('RANK', 0))
 
         defaults: DEFAULTS = {
             'lr': lr,