set fused to false to for pytorch nadamw

mlcommons · Sep 21, 2023 · d63214b · d63214b
1 parent 33a8a9f
commit d63214b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/baselines/nadamw/pytorch/submission.py b/baselines/nadamw/pytorch/submission.py
@@ -206,7 +206,9 @@ def init_optimizer_state(workload: spec.Workload,
               betas=(1.0 - hyperparameters.one_minus_beta1,
                      hyperparameters.beta2),
               eps=1e-8,
-              weight_decay=hyperparameters.weight_decay),
+              weight_decay=hyperparameters.weight_decay,
+              fused=False,
+          )
   }
 
   def pytorch_cosine_warmup(step_hint: int, hyperparameters, optimizer):

diff --git a/reference_algorithms/target_setting_algorithms/pytorch_nadamw.py b/reference_algorithms/target_setting_algorithms/pytorch_nadamw.py
@@ -203,7 +203,8 @@ def init_optimizer_state(workload: spec.Workload,
               lr=hyperparameters.learning_rate,
               betas=(hyperparameters.beta1, hyperparameters.beta2),
               eps=epsilon,
-              weight_decay=hyperparameters.weight_decay),
+              weight_decay=hyperparameters.weight_decay,
+              fused=False),
   }
 
   target_setting_step_hint = int(0.75 * workload.step_hint)