support DeepSpeed.

2026-10-29 09:07:10 +00:00 · 2024-12-24 10:59:21 +05:30
parent cf9be17819
commit 6329e9ed07
4 changed files with 140 additions and 52 deletions
@@ -15,6 +15,7 @@ class State:
    learning_rate: float = None
    train_batch_size: int = None
    generator: torch.Generator = None
+    num_update_steps_per_epoch: int = None

    # Hub state
    repo_id: str = None