Update app.py

app.py

@@ -51,6 +51,7 @@ class SmolLM2Config(PretrainedConfig):
 
 class SmolLM2ForCausalLM(PreTrainedModel):
     config_class = SmolLM2Config
+    _no_split_modules = ["LlamaDecoderLayer"]
     
     def __init__(self, config):
         super().__init__(config)
@@ -94,12 +95,17 @@ def load_model():
         }
         tokenizer.add_special_tokens(special_tokens)
         
+        # Load model without device_map
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
-            device_map="auto",
             pad_token_id=tokenizer.pad_token_id
         )
+        
+        # Move model to device manually
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = model.to(device)
+        
         # Resize token embeddings to match new tokenizer
         model.resize_token_embeddings(len(tokenizer))
         return model, tokenizer