actually use the registered buffer (#14)

teamtomo · Mar 27, 2023 · ff0bce5 · ff0bce5
1 parent adf8b8c
commit ff0bce5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/torch_cubic_spline_grids/_base_cubic_grid.py b/src/torch_cubic_spline_grids/_base_cubic_grid.py
@@ -38,7 +38,7 @@ def forward(self, u: torch.Tensor) -> torch.Tensor:
         u = self._coerce_to_batched_coordinates(u)  # (b, d)
         interpolated = [
             self._interpolation_function(
-                self._data, minibatch_u, matrix=self._interpolation_matrix
+                self._data, minibatch_u, matrix=self.interpolation_matrix
             )
             for minibatch_u in batch(u, n=self._minibatch_size)
         ]  # List[Tensor[(b, d)]]

diff --git a/tests/test_grids.py b/tests/test_grids.py
@@ -85,6 +85,13 @@ def test_calling_1d_grid_with_stacked_coords(grid_cls):
     assert torch.allclose(result, torch.tensor([0]).float())
 
 
+def test_interpolation_matrix_device():
+    model = CubicBSplineGrid1d(resolution=3)
+    assert model.interpolation_matrix.device == torch.device('cpu')
+    model.to(torch.device('meta'))
+    assert model.interpolation_matrix.device == torch.device('meta')
+
+
 @pytest.mark.parametrize(
     'grid_cls', [CubicBSplineGrid2d, CubicCatmullRomGrid2d]
 )