readme

src/main/inference.py

@@ -58,12 +58,12 @@ class Inference():
         self.image2image = ResUnetGenerator(input_nc=6,output_nc=3,num_downs=6,use_dropout=False)
 
         # Load pretrained weights to image2image model 
-        image2image_weight = torch.load(self.image2image_ckpt)['G']
+        image2image_weight = torch.load(self.image2image_ckpt, map_location=torch.device(device))['G']
         # Since the checkpoint of model was trained using DataParallel with multiple GPU
         # It required to wrap a model with DataParallel wrapper class
         self.image2image = DataParallel(self.image2image)
         # assgin weight to model
-        self.image2image.load_state_dict(image2image_weight, map_location=torch.device('cpu'))
+        self.image2image.load_state_dict(image2image_weight)