smollm2: minor fixes

Lightning-AI · Dec 16, 2024 · 1f1e737 · 1f1e737
1 parent 613eb77
commit 1f1e737
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/litgpt/prompts.py b/litgpt/prompts.py
@@ -332,7 +332,7 @@ def apply(self, prompt: str, **kwargs: str) -> str:
     "qwen2.5": Qwen2_5,
     "qwen2.5-math": Qwen2_5_Math,
     "qwq": QwQ,
-    "smollm2": SmolLM2, # SmolLM uses a different template
+    "smollm2": SmolLM2,
     "salamandra": Salamandra,
 }
 

diff --git a/tests/test_model.py b/tests/test_model.py
@@ -914,7 +914,7 @@ def test_against_original_salamandra(model_name, device, dtype):
 
 
 @torch.inference_mode()
-@pytest.mark.parametrize("model_name", ("SmolLM2-135M", "SmolLM2-1.7B"))
+@pytest.mark.parametrize("model_name", ("SmolLM2-135M", "SmolLM2-360M", "SmolLM2-1.7B"))
 @pytest.mark.parametrize(
     ("device", "dtype"),
     [