Merge pull request #115 from Jim137/Develop

Fix device compatibility issue in `KAN.initialize_from_another_model`
KindXiaoming · May 8, 2024 · cc507e3 · cc507e3
2 parents 70b7b8d + 96769e1
commit cc507e3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kan/KAN.py b/kan/KAN.py
@@ -190,10 +190,10 @@ def initialize_from_another_model(self, another_model, x):
         tensor(-0.0030)
         tensor(0.0506)
         '''
-        another_model(x)  # get activations
+        another_model(x.to(another_model.device))  # get activations
         batch = x.shape[0]
 
-        self.initialize_grid_from_another_model(another_model, x)
+        self.initialize_grid_from_another_model(another_model, x.to(another_model.device))
 
         for l in range(self.depth):
             spb = self.act_fun[l]