set foreach and fused false nadamw

mlcommons · Sep 22, 2023 · bf47c47 · bf47c47
1 parent 631d2d1
commit bf47c47
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/reference_algorithms/target_setting_algorithms/pytorch_nadamw.py b/reference_algorithms/target_setting_algorithms/pytorch_nadamw.py
@@ -58,6 +58,7 @@ def __init__(self,
       raise ValueError(f'Invalid weight_decay value: {weight_decay}')
     defaults = {
         'lr': lr, 'betas': betas, 'eps': eps, 'weight_decay': weight_decay,
+        'foreach': False, 'fused': False,
     }
     super().__init__(params, defaults)