Spaces:

akhaliq
/

BlendGAN

Runtime error

Ahsen Khaliq commited on Nov 19, 2021

Commit

4a4961c

1 Parent(s): 9e5c86d

Update psp_encoder/psp_encoders.py

Files changed (1) hide show

psp_encoder/psp_encoders.py CHANGED Viewed

@@ -141,7 +141,7 @@ class PSPEncoder(Module):
         print('Loading psp encoders weights from irse50!')
         encoder_ckpt = torch.load(encoder_ckpt_path, map_location='cpu')
         self.encoder.load_state_dict(get_keys(encoder_ckpt, 'encoder'), strict=True)
-        self.latent_avg = encoder_ckpt['latent_avg'].cuda()
         self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))

         print('Loading psp encoders weights from irse50!')
         encoder_ckpt = torch.load(encoder_ckpt_path, map_location='cpu')
         self.encoder.load_state_dict(get_keys(encoder_ckpt, 'encoder'), strict=True)
+        self.latent_avg = encoder_ckpt['latent_avg']
         self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))