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

tests/data/test_checkpoint.py

@@ -1,9 +1,10 @@
-import torch
 import tempfile
-import torch.distributed as dist
 
+import torch
+import torch.distributed as dist
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR
+
 from astrai.data.serialization import Checkpoint
 from astrai.parallel.setup import get_rank, spawn_parallel_fn
 

tests/data/test_dataset.py

@@ -1,8 +1,8 @@
-import torch
 import numpy as np
+import torch
 
-from astrai.data.serialization import save_h5
 from astrai.data.dataset import *
+from astrai.data.serialization import save_h5
 
 
 def test_dataset_loader_random_paths(base_test_env):

tests/data/test_sampler.py

@@ -1,5 +1,5 @@
-from astrai.trainer import *
 from astrai.data import *
+from astrai.trainer import *
 
 
 def test_random_sampler_consistency(random_dataset):