Update e4e/models/psp.py

e4e/models/psp.py

@@ -40,9 +40,11 @@ class pSp(nn.Module):
     def load_weights(self):
         if self.opts.checkpoint_path is not None:
             print('Loading e4e over the pSp framework from checkpoint: {}'.format(self.opts.checkpoint_path))
-            ckpt = torch.load(self.opts.checkpoint_path, map_location='cpu')
+            ckpt = torch.load(self.opts.checkpoint_path, map_location='cuda:0' if torch.cuda.is_available() else "cpu")
             self.encoder.load_state_dict(get_keys(ckpt, 'encoder'), strict=True)
+            self.encoder.to(self.device)
             self.decoder.load_state_dict(get_keys(ckpt, 'decoder'), strict=True)
+            self.decoder.to(self.device)
             self.__load_latent_avg(ckpt)
         else:
             print('Loading encoders weights from irse50!')