diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt
index 06e953ff..69e9d89d 100644
--- a/.ci/docker/requirements.txt
+++ b/.ci/docker/requirements.txt
@@ -6,3 +6,4 @@ sentencepiece
 tiktoken
 blobfile
 tabulate
+transformers
\ No newline at end of file
diff --git a/train_configs/chemlactica_125m.toml b/train_configs/chemlactica_125m.toml
index 5f303f55..842946c3 100644
--- a/train_configs/chemlactica_125m.toml
+++ b/train_configs/chemlactica_125m.toml
@@ -15,8 +15,10 @@ save_memory_snapshot_folder = "memory_snapshot"
 [metrics]
 log_freq = 1
 enable_color_printing = true
-enable_tensorboard = true
-save_tb_folder = "tb"
+enable_aim = true
+save_aim_folder = "aim"
+#aim_hash = "c6b4d8b340f74287b82ef928"
+#aim_experiment_name = "hello"
 
 [model]
 name = "opt"
diff --git a/train_configs/debug_model.toml b/train_configs/debug_model.toml
index 3eba5fef..fd2718d2 100644
--- a/train_configs/debug_model.toml
+++ b/train_configs/debug_model.toml
@@ -17,8 +17,8 @@ log_freq = 1
 enable_color_printing = true
 enable_aim = true
 save_aim_folder = "aim"
-aim_hash = "c6b4d8b340f74287b82ef928"
-aim_experiment_name = "hello"
+#aim_hash = "c6b4d8b340f74287b82ef928"
+#aim_experiment_name = "hello"
 
 [model]
 name = "llama3"
diff --git a/train_configs/galactica_125m_hf_to_titan.toml b/train_configs/galactica_125m_hf_to_titan.toml
index 1318d4cf..f6e0a314 100644
--- a/train_configs/galactica_125m_hf_to_titan.toml
+++ b/train_configs/galactica_125m_hf_to_titan.toml
@@ -15,8 +15,10 @@ save_memory_snapshot_folder = "memory_snapshot"
 [metrics]
 log_freq = 1
 enable_color_printing = true
-enable_tensorboard = true
-save_tb_folder = "tb"
+enable_aim = true
+save_aim_folder = "aim"
+#aim_hash = "c6b4d8b340f74287b82ef928"
+#aim_experiment_name = "hello"
 
 [model]
 name = "opt"