Update KANLayer.py

KindXiaoming · Jul 14, 2024 · 5e18f1b · 5e18f1b
1 parent d0ea020
commit 5e18f1b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kan/KANLayer.py b/kan/KANLayer.py
@@ -129,7 +129,7 @@ def __init__(self, in_dim=3, out_dim=2, num=5, k=3, noise_scale=0.1, scale_base=
         else:
             mask = 1.
 
-        scale_base = scale_base.to(device)
+        #scale_base = scale_base.to(device)
         self.scale_base = torch.nn.Parameter(torch.ones(in_dim, out_dim, device=device) * scale_base * mask).requires_grad_(sb_trainable)  # make scale trainable
         #else:
         #self.scale_base = torch.nn.Parameter(scale_base.to(device)).requires_grad_(sb_trainable)