Merge pull request #129 from Jim137/develop

Fix dtype error when input dtype is Double
KindXiaoming · May 8, 2024 · a666ec3 · a666ec3
2 parents fcefaba + d4305a3
commit a666ec3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/kan/spline.py b/kan/spline.py
@@ -97,6 +97,8 @@ def coef2curve(x_eval, grid, coef, k, device="cpu"):
     '''
     # x_eval: (size, batch), grid: (size, grid), coef: (size, coef)
     # coef: (size, coef), B_batch: (size, coef, batch), summer over coef
+    if coef.dtype != x_eval.dtype:
+        coef = coef.to(x_eval.dtype)
     y_eval = torch.einsum('ij,ijk->ik', coef, B_batch(x_eval, grid, k, device=device))
     return y_eval