diff --git a/reference_algorithms/development_algorithms/criteo1tb/criteo1tb_pytorch/submission.py b/reference_algorithms/development_algorithms/criteo1tb/criteo1tb_pytorch/submission.py
index d9d9c29b5..740193715 100644
--- 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(