Update app.py

app.py

@@ -229,7 +229,12 @@ current_directory = os.path.dirname(os.path.abspath(__file__))
 # Set the model path to the same directory as the Python file
 model_save_path = os.path.join(current_directory, 'model5k.pt')
 model = GPT(GPTConfig())
-model.load_state_dict(torch.load(model_save_path))
+if device == 'cpu':
+    # Load the model on CPU if no GPU is available
+    model.load_state_dict(torch.load(model_save_path, map_location=torch.device('cpu')))
+else:
+    # Load the model on GPU if available
+    model.load_state_dict(torch.load(model_save_path))
 model.to(device)
 model.eval()