make attention dropout static

Author: priyakasimbeg

algoperf/workloads/librispeech_conformer/librispeech_jax/models.py

@@ -428,6 +428,7 @@ def __call__(self, inputs, paddings, train, dropout_rate=DROPOUT_RATE):
       use_bias=True,
       broadcast_dropout=False,
       attention_fn=attention_fn,
+      dropout_rate=0.0,
       deterministic=not train,
     )(inputs_q=inputs, mask=attention_mask)
 

algoperf/workloads/wmt/wmt_jax/models.py

@@ -223,6 +223,7 @@ def __call__(self, inputs, encoder_mask=None, dropout_rate=DROPOUT_RATE):
       bias_init=cfg.bias_init,
       use_bias=False,
       broadcast_dropout=False,
+      dropout_rate=0.0,
       deterministic=cfg.deterministic,
     )(cfg.attention_temp * x, x, mask=encoder_mask)