Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -148,15 +148,7 @@ 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: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)
+    device = 'cuda'
     
     source_image = source_image.to(device)
     source_semantics = source_semantics.to(device)