set foreach false in nadamw

mlcommons · Sep 22, 2023 · 77c1d60 · 77c1d60
1 parent 29efbdd
commit 77c1d60
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/baselines/nadamw/pytorch/submission.py b/baselines/nadamw/pytorch/submission.py
@@ -48,7 +48,9 @@ def __init__(self,
                betas=(0.9, 0.999),
                eps=1e-8,
                weight_decay=1e-2,
-               fused=False,):
+               fused=False,
+               foreach=False,
+               ):
     if not 0.0 <= lr:
       raise ValueError(f'Invalid learning rate: {lr}')
     if not 0.0 <= eps:
@@ -61,7 +63,7 @@ def __init__(self,
       raise ValueError(f'Invalid weight_decay value: {weight_decay}')
     defaults = {
         'lr': lr, 'betas': betas, 'eps': eps, 'weight_decay': weight_decay,
-        'fused': fused,
+        'fused': fused, 'foreach': foreach,
     }
     super().__init__(params, defaults)
 

diff --git a/reference_algorithms/target_setting_algorithms/pytorch_nadamw.py b/reference_algorithms/target_setting_algorithms/pytorch_nadamw.py
@@ -45,7 +45,9 @@ def __init__(self,
                betas=(0.9, 0.999),
                eps=1e-8,
                weight_decay=1e-2,
-               fused=False):
+               fused=False,
+               foreach=False,
+               ):
     if not 0.0 <= lr:
       raise ValueError(f'Invalid learning rate: {lr}')
     if not 0.0 <= eps:
@@ -58,7 +60,7 @@ def __init__(self,
       raise ValueError(f'Invalid weight_decay value: {weight_decay}')
     defaults = {
         'lr': lr, 'betas': betas, 'eps': eps, 'weight_decay': weight_decay,
-        'fused': fused,
+        'fused': fused, 'foreach': foreach,
     }
     super().__init__(params, defaults)