Bug Fix for bert to run on entire validation dataset (#5493)

Signed-off-by: Shanmugam Ramasamy <[email protected]> Signed-off-by: Shanmugam Ramasamy <[email protected]>
NVIDIA · Nov 23, 2022 · b86bba8 · b86bba8
1 parent da953ae
commit b86bba8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -576,7 +576,7 @@ def validation_step(self, batch, batch_idx):
                 loss_sum = (
                     torch.vstack(loss_sum_tensors_list).sum(axis=0)
                     if len(loss_sum_tensors_list) > 0
-                    else torch.tensor(0.0)
+                    else torch.tensor([0.0, 0.0]).cuda()
                 )
                 return loss_sum
         else: