Fix grad accum type (#1296)

mosaicml · Jun 21, 2024 · e8ba9b7 · e8ba9b7
1 parent 13bd8f9
commit e8ba9b7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
@@ -184,7 +184,7 @@ class TrainConfig:
 
     # Fields created by `update_batch_size_info`
     n_gpus: int = MISSING
-    device_train_grad_accum: str = MISSING
+    device_train_grad_accum: Union[str, int] = MISSING
 
 
 TRAIN_CONFIG_KEYS = {field.name for field in fields(TrainConfig)}