Merge pull request #516 from taddyb/master

KindXiaoming · web-flow · commit ecde4ec3274d · 2025-01-19T13:49:39.000-05:00
Added support for saving checkpoints with CUDA device numbers
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

Original file line number	Diff line number	Diff line change
`@@ -534,6 +534,9 @@ def saveckpt(self, path='model'):`
`534`	`534`	`round = model.round,`
`535`	`535`	`device = str(model.device)`
`536`	`536`	`)`
	`537`	`+`
	`538`	`+ if dic["device"].isdigit():`
	`539`	`+ dic["device"] = int(model.device)`
`537`	`540`
`538`	`541`	`for i in range (model.depth):`
`539`	`542`	`dic[f'symbolic.funs_name.{i}'] = model.symbolic_fun[i].funs_name`