fix: 修复路径问题

khaosz/trainer/train_context.py

@@ -4,7 +4,7 @@ from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
 from khaosz.data import ResumableDistributedSampler
-from khaosz.data.checkpoint import Checkpoint
+from khaosz.data.serialization import Checkpoint
 from khaosz.trainer.strategy import StrategyFactory, BaseStrategy
 from khaosz.config.train_config import TrainConfig
 from khaosz.parallel.setup import get_current_device, get_world_size, get_rank

khaosz/trainer/trainer.py

@@ -10,7 +10,7 @@ from khaosz.trainer.train_callback import (
     SchedulerCallback
 )
 from khaosz.trainer.train_context import TrainContext, TrainContextBuilder
-from khaosz.data.checkpoint import Checkpoint
+from khaosz.data.serialization import Checkpoint
 from khaosz.parallel.setup import spawn_parallel_fn
 
 logger = logging.getLogger(__name__)