Upload app.py

app.py

@@ -42,7 +42,7 @@ def load_model():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = SmolLM(vocab_size, embed_dim, num_heads, num_layers, max_seq_len).to(device)
     
-    model.load_state_dict(torch.load(checkpoint_path))
+    model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
     #checkpoint = torch.load(checkpoint_path, map_location=device)
     #config = checkpoint['config']
     #model = GPT(config)