Upadated Inference.py

inference.py

@@ -97,8 +97,8 @@ def load_model(model_class, file_path):
     model = model_class(net_G=net_G)
     model.load_state_dict(torch.load(file_path, map_location=device))
 
-    resnet_weights = torch.load(file_path)
-    resnet_weights = torch.load("./model/res18-unet.pt")
+    resnet_weights = torch.load(file_path, map_location=device)
+    resnet_weights = torch.load("./model/res18-unet.pt", map_location=device)
     resnet_state_dict = resnet_weights['state_dict'] if 'state_dict' in resnet_weights else resnet_weights
 
     model_dict = model.state_dict()