Update warp_design_on_dress.py

warp_design_on_dress.py

@@ -19,8 +19,8 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
     dress_img = Image.open(dress_path).convert("RGB")
     design_img = Image.open(design_path).convert("RGB")
 
-    dress_tensor = tf(dress_img).unsqueeze(0).cuda()
-    design_tensor = tf(design_img).unsqueeze(0).cuda()
+    dress_tensor = tf(dress_img).unsqueeze(0).cpu()
+    design_tensor = tf(design_img).unsqueeze(0).cpu()
     design_mask = torch.ones_like(design_tensor[:, :1, :, :])  # full white mask
 
     # Fake agnostic: use the dress image itself