fix (#6529) (#6546)

Signed-off-by: Abhinav Khattar <[email protected]> Co-authored-by: Abhinav Khattar <[email protected]>
NVIDIA · May 26, 2023 · e4971ca · e4971ca
1 parent 5ad3d97
commit e4971ca
Showing 1 changed file with 1 addition and 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
@@ -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()