Upload app.py

app.py

@@ -46,7 +46,7 @@ def load_model():
     checkpoint = torch.load(checkpoint_path, map_location=device)
     #config = checkpoint['config']
     model = SmolLM(vocab_size, embed_dim, num_heads, num_layers, max_seq_len)
-    model.load_state_dict(checkpoint['model_state_dict'])
+    model.load_state_dict(checkpoint)
     model.to(device)
     model.eval()  # Set to evaluation mode