From 68e0cf4760d96efbd221d4fa17a2ccdaada4a283 Mon Sep 17 00:00:00 2001
From: Daniel King <daniel@mosaicml.com>
Date: Thu, 20 Jun 2024 22:41:52 -0700
Subject: [PATCH] fix

---
 llmfoundry/utils/config_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
index 52804f3520..79cdc225b2 100644
--- 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)}