Update src/audio2pose_models/audio_encoder.py

src/audio2pose_models/audio_encoder.py

@@ -19,7 +19,7 @@ class Conv2d(nn.Module):
         return self.act(out)
 
 class AudioEncoder(nn.Module):
-    def __init__(self, wav2lip_checkpoint):
+    def __init__(self, wav2lip_checkpoint, device='cpu'):
         super(AudioEncoder, self).__init__()
 
         self.audio_encoder = nn.Sequential(
@@ -42,7 +42,7 @@ class AudioEncoder(nn.Module):
             Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
 
         #### load the pre-trained audio_encoder\
-        wav2lip_state_dict = torch.load(wav2lip_checkpoint)['state_dict']
+        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=device)['state_dict']
         state_dict = self.audio_encoder.state_dict()
 
         for k,v in wav2lip_state_dict.items():