Forced model to load to CPU (map_location)

app.py

@@ -25,12 +25,13 @@ API_KEY = os.getenv("HUGGINGFACE_API_KEY")
 API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
 
 # --- LOAD DATA AND MODEL ---
+# map_location forces the model to be loaded on the CPU for huggingface compatibility
 movies_df = pd.read_csv("./sampled_movie_dataset/movies_metadata.csv")  # Load your movie data
-data = torch.load("./PyGdata.pt")
+data = torch.load("./PyGdata.pt", map_location=device('cpu'))
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = model_def.Model(hidden_channels=32).to(device) 
-model.load_state_dict(torch.load("PyGTrainedModelState.pt"))
+model.load_state_dict(torch.load("PyGTrainedModelState.pt", map_location=device('cpu')))
 model.eval()
 
 # --- STREAMLIT APP ---