diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py
index 946df3da2aa5..c80d2272613e 100644
--- a/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py
+++ b/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+from functools import partial
 from typing import Any, Optional
 
 import torch