fix: get_parameters

kozistr · Dec 3, 2024 · 64486e4 · 64486e4
1 parent a043fa5
commit 64486e4
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/pytorch_optimizer/optimizer/muon.py b/pytorch_optimizer/optimizer/muon.py
@@ -96,8 +96,6 @@ def __init__(
         params = self.get_parameters(params)
         adamw_params = self.get_parameters(adamw_params) if adamw_params is not None else []
         params.extend(adamw_params)
-        # print(params)
-        # print(adamw_params)
 
         self.world_size: int = os.environ.get('WORLD_SIZE', 1)
         self.rank: int = os.environ.get('RANK', 0)
@@ -128,7 +126,7 @@ def get_parameters(params: PARAMETERS) -> PARAMETERS:
         new_params = []
         for group in params:
             if isinstance(group, dict) and 'params' in group:
-                new_params.extend(group['params'])
+                new_params.extend(list(group['params']))
             else:
                 new_params.append(group)