Fix cuda/cpu issues

jloveric · Jun 19, 2024 · 9ee999c · 9ee999c
1 parent b3f86f2
commit 9ee999c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/config/block_mnist.yaml b/config/block_mnist.yaml
@@ -1,6 +1,6 @@
 max_epochs: 1
-accelerator: 'cpu'
-n: 5
+accelerator: 'cuda'
+n: 10
 batch_size: 16
 layer_type: polynomial_3d
 train_fraction: 1.0

diff --git a/high_order_layers_torch/LagrangePolynomial.py b/high_order_layers_torch/LagrangePolynomial.py
@@ -82,7 +82,7 @@ def __init__(
     ):
         self.n = n
         self.dimensions = dimensions
-        self.X = (length / 2.0) * chebyshevLobatto(n)
+        self.X = (length / 2.0) * chebyshevLobatto(n).to(device)
         self.device = device
         self.denominators = self._compute_denominators()
         self.num_basis = int(math.pow(n, dimensions))
@@ -102,14 +102,13 @@ def __call__(self, x, index: list[int]):
         :param index : [dimensions]
         :returns: basis value [batch, inputs]
         """
-
         x_diff = x.unsqueeze(-1) - self.X  # [batch, inputs, dimensions, basis]
         r = 1.0
         for i, basis_i in enumerate(index):
             b = torch.where(
-                torch.arange(self.n) != basis_i,
+                torch.arange(self.n, device=self.device) != basis_i,
                 x_diff[:, :, i, :] / self.denominators[basis_i],
-                torch.tensor(1.0),
+                torch.tensor(1.0, device=self.device),
             )
             r *= torch.prod(b, dim=-1)