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

…lining (NVIDIA#6449) Co-authored-by: ANMOL GUPTA <[email protected]> Co-authored-by: Eric Harper <[email protected]>
anmolgupt · May 12, 2023 · 0bfd158 · 0bfd158
1 parent d9650a5
commit 0bfd158
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
@@ -647,6 +647,10 @@ def validation_step(self, dataloader_iter, batch_idx):
             The list of microbatches is then piped through the pipeline using Apex 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
@@ -665,6 +669,10 @@ def validation_step(self, dataloader_iter, batch_idx):
             batch_p2p_comm=self.cfg.get('batch_p2p_comm', True),
         )
 
+        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):