fix tests

Author: SkafteNicki

stochman/curves.py

@@ -13,7 +13,6 @@ def __init__(
         end: torch.Tensor,
         num_nodes: int = 5,
         requires_grad: bool = True,
-        device=None,
         *args,
         **kwargs,
     ) -> None:
@@ -201,7 +200,7 @@ def forward(self, t: torch.Tensor) -> torch.Tensor:
             (torch.floor(tt * num_edges).clamp(min=0, max=num_edges - 1).long())  # Bx|t|
             .unsqueeze(2)
             .repeat(1, 1, D)
-        )  # Bx|t|xD, this assumes that nodes are equi-distant
+        ).to(self.device)  # Bx|t|xD, this assumes that nodes are equi-distant
         result = torch.gather(a, 1, idx) * tt.unsqueeze(2) + torch.gather(b, 1, idx)  # Bx|t|xD
         if B == 1:
             result = result.squeeze(0)  # |t|xD
@@ -221,7 +220,7 @@ def __getitem__(self, indices: int) -> "DiscreteCurve":
         return C
 
     def __setitem__(self, indices, curves) -> None:
-        self.params[indices] = curves.params.squeeze()
+        self.params[indices].data = curves.params.squeeze()
 
     def constant_speed(
         self, metric=None, t: Optional[torch.Tensor] = None
@@ -399,7 +398,7 @@ def __getitem__(self, indices: int) -> "CubicSpline":
         return C
 
     def __setitem__(self, indices, curves) -> None:
-        self.params[indices] = curves.params
+        self.params[indices].data = curves.params
 
     def deriv(self, t: Optional[torch.Tensor] = None) -> torch.Tensor:
         """

tests/test_curves.py

@@ -8,7 +8,7 @@
 class TestCurves:
     @pytest.mark.parametrize("requires_grad", [True, False])
     @pytest.mark.parametrize("batch_dim", [1, 5])
-    @pytest.mark.parametrize("device", ["cpu", "cuda"])
+    @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
     def test_curve_evaluation(self, curve_class, requires_grad, batch_dim, device):
         if not torch.cuda.is_available() and device == "cuda":
             pytest.skip("test requires cuda")
@@ -27,7 +27,7 @@ def test_curve_evaluation(self, curve_class, requires_grad, batch_dim, device):
         assert c.params.device == torch.device(device)
 
         eval_nodes = 10
-        t = torch.linspace(0, 1, eval_nodes)
+        t = torch.linspace(0, 1, eval_nodes).to(device)
         out = c(t)
         assert isinstance(out, torch.Tensor)
         assert out.device == torch.device(device)