Update app.py

app.py

@@ -18,9 +18,9 @@ model = SmollM(config['model']['model_config'])
 tokenizer = AutoTokenizer.from_pretrained(config['tokenizer']['tokenizer_name_or_path'])
 
 # Load your custom model (adjust as necessary for your model's implementation)
-model_path = "model.pth"  # Replace with the path to your model weights
-checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
-model.load_state_dict(checkpoint['model_state_dict'])
+model_path = "model.pt"  # Replace with the path to your model weights
+checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
+model.load_state_dict(checkpoint)
 model.eval()  # Set the model to evaluation mode
 
 def generate_tokens(model, tokenizer, prompt, max_length=50, device="cuda"):