Merge pull request #182 from eigenvivek/fix-compile

Move dims to device for torch.compile

Author: eigenvivek

diffdrr/siddon.py

@@ -15,7 +15,7 @@ def siddon_raycast(
     eps: float = 1e-8,
 ):
     """An auto-differentiable implementation of the raycasting algorithm known as Siddon's method."""
-    dims = torch.tensor(volume.shape) + 1
+    dims = torch.tensor(volume.shape).to(source) + 1
     alphas, maxidx = _get_alphas(source, target, spacing, dims, eps)
     alphamid = (alphas[..., 0:-1] + alphas[..., 1:]) / 2
     voxels = _get_voxel(alphamid, source, target, volume, spacing, dims, maxidx, eps)

notebooks/api/01_siddon.ipynb

@@ -91,7 +91,7 @@
     "    eps: float=1e-8,\n",
     "):\n",
     "    \"\"\"An auto-differentiable implementation of the raycasting algorithm known as Siddon's method.\"\"\"\n",
-    "    dims = torch.tensor(volume.shape) + 1\n",
+    "    dims = torch.tensor(volume.shape).to(source) + 1\n",
     "    alphas, maxidx = _get_alphas(source, target, spacing, dims, eps)\n",
     "    alphamid = (alphas[..., 0:-1] + alphas[..., 1:]) / 2\n",
     "    voxels = _get_voxel(alphamid, source, target, volume, spacing, dims, maxidx, eps)\n",