minor bug fix (NVIDIA#8696)

Signed-off-by: arendu <[email protected]>
anmolgupt · Mar 19, 2024 · 7332754 · 7332754
1 parent 1810f1f
commit 7332754
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/information_retrieval/megatron_gpt_embedding_model.py b/nemo/collections/nlp/models/information_retrieval/megatron_gpt_embedding_model.py
@@ -241,7 +241,7 @@ def inference_step_validation_call(self, batch, batch_idx, data_cfg, dataloader_
         }
         return outputs
 
-    def gather_and_maybe_write_predictions(self, output, data_cfg, mode, dataloader_idx=0):
+    def gather_and_maybe_write_predictions(self, output, data_cfg, mode, averaged_metric, dataloader_idx=0):
         if not data_cfg.get("write_embeddings_to_file", False):
             return True
         gathered_output_batches = [None for _ in range(parallel_state.get_data_parallel_world_size())]