Still figuring out this CPU thing.

app.py

@@ -30,9 +30,9 @@ movies_df = pd.read_csv("./sampled_movie_dataset/movies_metadata.csv")  # Load y
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
-data = torch.load("./PyGdata.pt", map_location=device('cpu'))
+data = torch.load("./PyGdata.pt", map_location=device)
 model = model_def.Model(hidden_channels=32).to(device) 
-model.load_state_dict(torch.load("PyGTrainedModelState.pt", map_location=device('cpu')))
+model.load_state_dict(torch.load("PyGTrainedModelState.pt", map_location=device))
 model.eval()
 
 # --- STREAMLIT APP ---

model_def.py

@@ -6,7 +6,7 @@ from torch_geometric.nn import SAGEConv, to_hetero, Linear
 from dotenv import load_dotenv
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-data = torch.load("./PyGdata.pt", map_location=device('cpu'))
+data = torch.load("./PyGdata.pt", map_location=device)
 
 class GNNEncoder(torch.nn.Module):
     def __init__(self, hidden_channels, out_channels):

visualizer.py

@@ -19,7 +19,7 @@ import yaml
 
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-data = torch.load("./PyGdata.pt", map_location=device('cpu'))
+data = torch.load("./PyGdata.pt", map_location=device)
 
 
 movies_df = pd.read_csv("./sampled_movie_dataset/movies_metadata.csv")
@@ -65,7 +65,7 @@ class Model(torch.nn.Module):
     
 model = Model(hidden_channels=32).to(device)
 model2 = Model(hidden_channels=32).to(device)
-model.load_state_dict(torch.load("PyGTrainedModelState.pt"), map_location=device('cpu'))
+model.load_state_dict(torch.load("PyGTrainedModelState.pt"), map_location=device)
 model.eval()
 
 total_users = data['user'].num_nodes