Set deterministic=True for inference.

google · Nov 1, 2023 · 1612dc7 · 1612dc7
1 parent 9d6b6db
commit 1612dc7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/paxml/contrib/gpu/scripts_gpu/te_helper.py b/paxml/contrib/gpu/scripts_gpu/te_helper.py
@@ -141,7 +141,8 @@ def __call__(self,
             inputs=x_in,
             attention_mask=attention_mask,
             encoded=cross_inputs,
-            encoder_decoder_mask=cross_attention_mask)
+            encoder_decoder_mask=cross_attention_mask,
+            deterministic=self.do_eval)
         x_out = checkpoint_name(x_out, 'transformer_layer_out')
     return x_out