update config names in tests

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
NVIDIA · Sep 27, 2024 · b7c0bfa · b7c0bfa
1 parent b57fd6c
commit b7c0bfa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/collections/llm/gpt/model/test_mistral.py b/tests/collections/llm/gpt/model/test_mistral.py
@@ -1,6 +1,6 @@
 import torch.nn.functional as F
 
-from nemo.collections.llm.gpt.model.mistral import MistralConfig7B, MistralNeMo2407Config12B, MistralNeMo2407Config123B
+from nemo.collections.llm.gpt.model.mistral import MistralConfig7B, MistralNeMoConfig12B, MistralNeMoConfig123B
 
 
 def test_mistral_config7b():
@@ -25,7 +25,7 @@ def test_mistral_config7b():
 
 
 def test_mistral_nemo_config_12b():
-    config = MistralNeMo2407Config12B()
+    config = MistralNeMoConfig12B()
     assert config.normalization == "RMSNorm"
     assert config.activation_func == F.silu
     assert config.position_embedding_type == "rope"
@@ -49,7 +49,7 @@ def test_mistral_nemo_config_12b():
 
 
 def test_mistral_nemo_config_123b():
-    config = MistralNeMo2407Config123B()
+    config = MistralNeMoConfig123B()
     assert config.normalization == "RMSNorm"
     assert config.activation_func == F.silu
     assert config.position_embedding_type == "rope"