Should be mlguess.torch.checkpoint

ai2es · Jul 20, 2024 · 368c0fb · 368c0fb
1 parent 469eda9
commit 368c0fb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mlguess/torch/trainer.py b/mlguess/torch/trainer.py
@@ -12,8 +12,8 @@
 import tqdm
 from torch.cuda.amp import autocast
 import optuna
-from credit.models.checkpoint import TorchFSDPCheckpointIO
-
+#from credit.models.checkpoint import TorchFSDPCheckpointIO
+from mlguess.torch.checkpoint import TorchFSDPCheckpointIO
 
 def cleanup():
     dist.destroy_process_group()