update

drag_gan.py

@@ -168,7 +168,7 @@ def drag_gan(g_ema, latent: torch.Tensor, noise, F, handle_points, target_points
         points = []
         for i in range(x - d, x + d):
             for j in range(y - d, y + d):
-                points.append(torch.tensor([i, j]).float().cuda())
+                points.append(torch.tensor([i, j]).float().to(latent.device))
         return points
 
     F0 = F.detach().clone()
@@ -213,11 +213,7 @@ def drag_gan(g_ema, latent: torch.Tensor, noise, F, handle_points, target_points
                 miny = 1e9
                 for qi in neighbor(int(handle_points[i][0]), int(handle_points[i][1]), r2):
                     # f2 = F2[..., int(qi[0]), int(qi[1])]
-                    try:
-                        f2 = bilinear_interpolate_torch(F2, qi[0], qi[1])
-                    except:
-                        import ipdb
-                        ipdb.set_trace()
+                    f2 = bilinear_interpolate_torch(F2, qi[0], qi[1])
                     v = torch.norm(f2 - f0, p=1)
                     if v < minv:
                         minv = v