harlanhong / CVPR2022-DaGAN

Official code for CVPR2022 paper: Depth-Aware Generative Adversarial Network for Talking Head Video Generation
https://harlanhong.github.io/publications/dagan.html
Other
957 stars 125 forks source link

Add torch.device('cpu') when loading encoder/decoder weights for CPU use #51

Closed tikitong closed 1 year ago

tikitong commented 1 year ago

In my case, inference with cpu via the --cpu argument needs the addition of map_location=torch.device('cpu') in torch.load() to succeed. Can you confirm? Many thanks!

https://github.com/harlanhong/CVPR2022-DaGAN/blob/78b22edcdbb4192b81c5adf343f980b42cddfe5d/demo.py#L166

if opt.cpu:
  loaded_dict_enc = torch.load('checkpoints/depth_face_model/encoder.pth', map_location=torch.device('cpu'))
  loaded_dict_dec = torch.load('checkpoints/depth_face_model/depth.pth', map_location=torch.device('cpu'))
harlanhong commented 1 year ago

Correct. Thanks you.