fix cuda bug in undistort_image

Spiideo · Dec 5, 2024 · 9339821 · 9339821
1 parent 64cdec5
commit 9339821
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/sskit/camera.py b/sskit/camera.py
@@ -89,7 +89,7 @@ def project_on_ground(camera_matrix, dist_poly, image, width=70, height=120, res
 
 def undistort_image(dist_poly, image, zoom=1):
     h, w = image.shape[-2:]
-    grid = (grid2d(w, h) - torch.tensor([w/2, h/2])) / w / zoom
+    grid = (grid2d(w, h) - torch.tensor([w/2, h/2])).to(image.device) / w / zoom
     dgrid = distort(dist_poly, grid)
     return sample_image(image, dgrid[None])
 

diff --git a/sskit/version.py b/sskit/version.py
@@ -1,2 +1,2 @@
-__version_info__ = (0, 2, 0)
+__version_info__ = (0, 2, 1)
 __version__ = '.'.join(str(i) for i in __version_info__)