Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -148,10 +148,15 @@ def make_animation(source_image, source_semantics, target_semantics,
                    yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None,
                    use_exp=True):
 
-    device='cuda'
-    generator = torch.nn.DataParallel(generator).to(device)
-    kp_detector = torch.nn.DataParallel(kp_detector).to(device)
-    mapping = torch.nn.DataParallel(mapping).to(device)
+    device = 'cuda:0'
+    generator = generator.to(device)
+    kp_detector = kp_detector.to(device)
+    mapping = mapping.to(device)
+
+    # Wrap the models in DataParallel to use all available GPUs
+    generator = torch.nn.DataParallel(generator)
+    kp_detector = torch.nn.DataParallel(kp_detector)
+    mapping = torch.nn.DataParallel(mapping)
     
     source_image = source_image.to(device)
     source_semantics = source_semantics.to(device)