Merge pull request #516 from taddyb/master

Added support for saving checkpoints with CUDA device numbers
KindXiaoming · Jan 19, 2025 · ecde4ec · ecde4ec
2 parents 0a452a0 + 1625986
commit ecde4ec
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/kan/MultKAN.py b/kan/MultKAN.py
@@ -534,6 +534,9 @@ def saveckpt(self, path='model'):
             round = model.round,
             device = str(model.device)
         )
+
+        if dic["device"].isdigit():
+            dic["device"] = int(model.device)
 
         for i in range (model.depth):
             dic[f'symbolic.funs_name.{i}'] = model.symbolic_fun[i].funs_name