refactor: 拆分engine.py 文件

2026-04-05 00:07:21 +08:00
parent 861d33b1a1
commit 2b26f03bd3
4 changed files with 479 additions and 553 deletions
@@ -4,6 +4,7 @@ import torch

 from astrai.config.param_config import ModelParameter

+
 def test_model_parameter(test_env):
    save_dir = os.path.join(test_env["test_dir"], "save")
    model_param = ModelParameter(
@@ -30,4 +31,4 @@ def test_transformer(test_env):
        test_env["transformer_config"].max_len,
        test_env["transformer_config"].vocab_size,
    )
-    assert output_logits.shape == target_shape
+    assert output_logits.shape == target_shape