Don't pass kwargs that aren't accepted by model_fn

Signed-off-by: Edward Z. Yang <[email protected]>
mlcommons · Aug 10, 2023 · 716cb64 · 716cb64
1 parent 5577b32
commit 716cb64
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/reference_algorithms/development_algorithms/criteo1tb/criteo1tb_pytorch/submission.py b/reference_algorithms/development_algorithms/criteo1tb/criteo1tb_pytorch/submission.py
@@ -74,8 +74,6 @@ def update_params(workload: spec.Workload,
       model_state=model_state,
       mode=spec.ForwardPassMode.TRAIN,
       rng=rng,
-      dropout_rate=None,
-      aux_dropout_rate=None,
       update_batch_norm=False)
 
   loss_dict = workload.loss_fn(