Proposed WAR for gpt3 eval hang with PP (NVIDIA#7927)

Signed-off-by: yaoyu-33 <[email protected]> Co-authored-by: Eric Harper <[email protected]>
rohitrango · Dec 1, 2023 · 93b2c7f · 93b2c7f
1 parent 44c7928
commit 93b2c7f
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/nemo/collections/nlp/modules/common/text_generation_utils.py b/nemo/collections/nlp/modules/common/text_generation_utils.py
@@ -507,12 +507,8 @@ def synced_generate(
 
             if compute_logprob:
                 precision = model._trainer.precision
-                if precision in [16, "16"]:
-                    dtype = torch.float16
-                elif precision in ['bf16', 'bf16-mixed']:
-                    dtype = torch.bfloat16
-                else:
-                    dtype = torch.float32
+                dtype = torch.float32
+
                 output_logits = torch.empty(
                     tokens.size(0), context_length - 1, dtype=dtype, device=torch.device("cuda")
                 )