From 368c0fbbdd5e4c0d99f2377c6fbe855351dae1fa Mon Sep 17 00:00:00 2001
From: Israt Jahan <ijahan@casper-login1.hpc.ucar.edu>
Date: Fri, 19 Jul 2024 20:29:17 -0600
Subject: [PATCH] Should be mlguess.torch.checkpoint

---
 mlguess/torch/trainer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlguess/torch/trainer.py b/mlguess/torch/trainer.py
index 2e3b3d2..541296c 100644
--- 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()