Update logic2.py

logic2.py

@@ -67,10 +67,10 @@ def load_hetero_data():
         data = pickle.load(file)
         return data
 
-def load_model(train_data, val_data, test_data):
+def load_model(data):
     model = Model(hidden_channels=32)
     with torch.no_grad():
-        model.encoder(train_data.x_dict, train_data.edge_index_dict)
+        model.encoder(data.x_dict, data.edge_index_dict)
     model.load_state_dict(torch.load('model.pt',map_location=torch.device('cpu')))
     model.eval()
     return model