refactor: 优化参数传递，清理导入样式

astrai/trainer/train_context.py

@@ -1,16 +1,16 @@
+from dataclasses import dataclass, field
+from typing import Optional, Self
+
 import torch.nn as nn
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
+from astrai.config.train_config import TrainConfig
 from astrai.data import ResumableDistributedSampler
 from astrai.data.serialization import Checkpoint
-from astrai.trainer.strategy import StrategyFactory, BaseStrategy
-from astrai.config.train_config import TrainConfig
-from astrai.parallel.setup import get_current_device, get_world_size, get_rank
-
-from dataclasses import dataclass, field
-from typing import Optional, Self
+from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
 
 @dataclass