Update mymodels.py

mymodels.py

@@ -258,8 +258,8 @@ class Color2Sketch(nn.Module):
         self.decoder = Decoder()
         if pretrained:
             print('Loading pretrained {0} model...'.format('Color2Sketch'), end=' ')
-            checkpoint = torch.load('color2edge.pth')
-            self.load_state_dict(checkpoint['netG'], strict=False)
+            checkpoint = torch.load('color2edge.pth', map_location = "cuda" if torch.cuda.is_available() else "cpu")
+            self.load_state_dict(checkpoint['netG'], strict=True)
             print("Done!")
         else:
             self.apply(weights_init)
@@ -395,7 +395,7 @@ class Sketch2Color(nn.Module):
         self.decoder = Decoder()
         if pretrained:
             print('Loading pretrained {0} model...'.format('Sketch2Color'), end=' ')
-            checkpoint = torch.load('edge2color.pth')
+            checkpoint = torch.load('edge2color.pth', map_location = "cuda" if torch.cuda.is_available() else "cpu")
             self.load_state_dict(checkpoint['netG'], strict=True)
             print("Done!")
         else: