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

Convert scale to tensor in appropriate device

Author: ErnstRoell

dect/ect.py

@@ -78,6 +78,10 @@ def compute_ect(
     """
 
     # ecc.shape[0], index.max().item() + 1, ecc.shape[2],
+
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:
@@ -165,6 +169,10 @@ def compute_ect_point_cloud(
         point clouds (thus ECT's), N is the number of direction and R is the
         resolution.
     """
+
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     lin = torch.linspace(
         start=-radius, end=radius, steps=resolution, device=x.device
     ).view(-1, 1, 1)
@@ -208,6 +216,9 @@ def compute_ect_points(
         The index tensor is assumed to start at 0.
     """
 
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:
@@ -273,6 +284,9 @@ def compute_ect_edges(
         The index tensor is assumed to start at 0.
     """
 
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:
@@ -357,6 +371,9 @@ def compute_ect_mesh(
         The index tensor is assumed to start at 0.
     """
 
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else: