state.py

from dataclasses import dataclass

@dataclass
class ModelState:
    vocab_size: int
    num_layers: int
    hidden_size: int
    intermediate_size: int