fix to set model in the eval mode during validation with virtual pipe…

…lining (#6449) Co-authored-by: ANMOL GUPTA <[email protected]> Co-authored-by: Eric Harper <[email protected]>
NVIDIA · Apr 19, 2023 · a80467c · a80467c
1 parent a365879
commit a80467c
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -632,6 +632,10 @@ def validation_step(self, dataloader_iter, batch_idx):
             The list of microbatches is then piped through the pipeline using megatron-core fwd/bwd functions.
         """
 
+        if isinstance(self.model, list):
+            for model_module in self.model:
+                model_module.eval()
+
         tensor_shape = [self.cfg.encoder_seq_length, self.cfg.micro_batch_size, self.cfg.hidden_size]
 
         # run forward passes for an entire global batch
@@ -649,6 +653,10 @@ def validation_step(self, dataloader_iter, batch_idx):
             sequence_parallel=self.cfg.get('sequence_parallel', False),
         )
 
+        if isinstance(self.model, list):
+            for model_module in self.model:
+                model_module.train()
+
         # only the last stage of the pipeline returns losses
         if losses_reduced_per_micro_batch:
             if self.cfg.data.get('validation_drop_last', True):