Merge pull request #10 from aidos-lab/dev-jpga

Add device for linspace in function compute_ect_edges

Author: ErnstRoell

dect/ect.py

@@ -287,7 +287,7 @@ def compute_ect_edges(
 
     # Node heights have shape [num_points, num_directions]
     nh = x @ v
-    lin = torch.linspace(-radius, radius, resolution).view(-1, 1, 1)
+    lin = torch.linspace(-radius, radius, resolution, device=x.device).view(-1, 1, 1)
     ecc = torch.nn.functional.sigmoid(scale * torch.sub(lin, nh))
     output = torch.zeros(
         size=out_shape,