From 756d7eaf94cf00f57ef1197f5a0fb8188c3d6c46 Mon Sep 17 00:00:00 2001 From: kozistr Date: Tue, 3 Dec 2024 22:20:28 +0900 Subject: [PATCH] fix: type --- pytorch_optimizer/optimizer/muon.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytorch_optimizer/optimizer/muon.py b/pytorch_optimizer/optimizer/muon.py index 2a4abf47..7772949b 100644 --- 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,