fix criteo1tb

Author: priyakasimbeg

algoperf/workloads/criteo1tb/criteo1tb_jax/models.py

@@ -223,11 +223,7 @@ def scaled_init(key, shape, dtype=jnp.float_):
         top_mlp_input = nn.relu(top_mlp_input)
         if self.use_layer_norm:
           top_mlp_input = nn.LayerNorm()(top_mlp_input)
-      if (
-        dropout_rate is not None
-        and dropout_rate > 0.0
-        and layer_idx == num_layers_top - 2
-      ):
+      if dropout_rate is not None and layer_idx == num_layers_top - 2:
         top_mlp_input = Dropout(dropout_rate, deterministic=not train)(
           top_mlp_input, rate=dropout_rate
         )

algoperf/workloads/imagenet_resnet/imagenet_jax/models.py

@@ -89,10 +89,8 @@ def __call__(
     x: spec.Tensor,
     update_batch_norm: bool = True,
     use_running_average_bn: Optional[bool] = None,
-    dropout_rate: float = 0.0,
   ) -> spec.Tensor:
     conv = functools.partial(nn.Conv, use_bias=False, dtype=self.dtype)
-    del dropout_rate  # unused
     # Preserve default behavior for backwards compatibility
     if use_running_average_bn is None:
       use_running_average_bn = not update_batch_norm

algoperf/workloads/imagenet_resnet/imagenet_jax/workload.py

@@ -154,9 +154,11 @@ def model_fn(
     rng: spec.RandomState,
     update_batch_norm: bool,
     use_running_average_bn: Optional[bool] = None,
+    dropout_rate: Optional[float] = None,
   ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     del mode
     del rng
+    del dropout_rate
     variables = {'params': params, **model_state}
     if update_batch_norm:
       logits, new_model_state = self._model.apply(