move import to avoid te issue

NVIDIA · Oct 29, 2024 · 213ffca · 213ffca
1 parent cf33e03
commit 213ffca
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/collections/llm/gpt/model/gemma2.py b/nemo/collections/llm/gpt/model/gemma2.py
@@ -24,7 +24,6 @@
 from nemo.collections.llm.fn.activation import openai_gelu
 from nemo.collections.llm.gpt.model.base import GPTConfig, GPTModel
 from nemo.collections.llm.utils import Config
-from nemo.collections.nlp.models.language_modeling.megatron.gemma2.gemma2_spec import get_gemma2_layer_spec
 from nemo.lightning import OptimizerModule, io, teardown
 from nemo.lightning.pytorch.utils import dtype_from_hf
 
@@ -36,6 +35,8 @@
 
 
 def gemma2_layer_spec(config: "GPTConfig") -> ModuleSpec:
+    from nemo.collections.nlp.models.language_modeling.megatron.gemma2.gemma2_spec import get_gemma2_layer_spec
+
     return get_gemma2_layer_spec()