From 8b5a1bbddad8a5bf7f35e465e81bdcbbf94aede4 Mon Sep 17 00:00:00 2001
From: Daniel King <43149077+dakinggg@users.noreply.github.com>
Date: Sat, 22 Jun 2024 19:10:27 -0700
Subject: [PATCH] Add `all` transforms to train script (#1300)

---
 scripts/train/train.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/scripts/train/train.py b/scripts/train/train.py
index 84df79722a..4620254135 100644
--- a/scripts/train/train.py
+++ b/scripts/train/train.py
@@ -50,7 +50,6 @@
     make_dataclass_and_log_config,
     pop_config,
     process_init_device,
-    update_batch_size_info,
 )
 from llmfoundry.utils.exceptions import (
     BaseContextualError,
@@ -197,7 +196,7 @@ def main(cfg: DictConfig) -> Trainer:
         cfg,
         TrainConfig,
         TRAIN_CONFIG_KEYS,
-        transforms=[update_batch_size_info],
+        transforms='all',
     )
 
     # Set logging level