From 4d3f50ac2ff835f1d7548ffa7448713c982a48dc Mon Sep 17 00:00:00 2001 From: Abhinav Khattar Date: Mon, 1 May 2023 18:59:54 -0700 Subject: [PATCH] fix Signed-off-by: Abhinav Khattar --- .../nlp/models/language_modeling/megatron_finetune_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/collections/nlp/models/language_modeling/megatron_finetune_model.py b/nemo/collections/nlp/models/language_modeling/megatron_finetune_model.py index c4cfcfdad1ff..fb58ec6a843b 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron_finetune_model.py +++ b/nemo/collections/nlp/models/language_modeling/megatron_finetune_model.py @@ -284,7 +284,7 @@ def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only): _, seq_length = batch[0].shape _, dec_seq_length = batch[1].shape - tensor_shape = [seq_length, get_micro_batch_size(), self.hidden_size] + tensor_shape = [seq_length, get_micro_batch_size(), self.cfg.encoder.hidden_size] data_iter = get_iterator_k_split(batch, get_num_microbatches()) fwd_bwd_function = get_forward_backward_func()