fixed torch.load

app.py

@@ -169,7 +169,7 @@ model = BigramLanguageModel()
 checkpoint_path = 'checkpoint.pth'
 
 # Load the model checkpoint
-checkpoint = torch.load(checkpoint_path)
+checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
 model.load_state_dict(checkpoint['model_state_dict'])
 model.eval()
 model.to(device)