refactor: 修改 StepMonitorCallback, 分离职责

2026-03-04 19:45:39 +08:00
parent b53e10aac4
commit 5713b55500
5 changed files with 48 additions and 57 deletions
@@ -4,19 +4,14 @@ import numpy as np
 import tempfile
 import shutil
 import torch
-
 import pytest
-import matplotlib
-from torch.utils.data import Dataset

+from torch.utils.data import Dataset
 from khaosz.config.model_config import ModelConfig
 from khaosz.data.tokenizer import BpeTokenizer
 from khaosz.model.transformer import Transformer


-matplotlib.use("Agg")
-
-
 class RandomDataset(Dataset):
    def __init__(self, length=None, max_length=64, vocab_size=1000):
        self.length = length or int(np.random.randint(100, 200))