From 73aa406b966b3dadf20aaeef7e3c0c3e5937eeb9 Mon Sep 17 00:00:00 2001
From: Xuanlei Zhao <xuanlei.zhao@gmail.com>
Date: Thu, 28 Dec 2023 15:48:04 +0800
Subject: [PATCH] update

---
 applications/ColossalMoE/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/applications/ColossalMoE/train.py b/applications/ColossalMoE/train.py
index d0f950e95..79052f910 100644
--- a/applications/ColossalMoE/train.py
+++ b/applications/ColossalMoE/train.py
@@ -224,7 +224,7 @@ def main():
         with tqdm(
             range(total_len),
             desc=f"Epoch [{epoch + 1}/{args.num_epoch}]",
-            disable=not coordinator.is_master(),
+            disable=not coordinator.is_master() if use_pipeline == False else not is_pp_last_stage,
         ) as pbar:
             for step in pbar:
                 if use_pipeline: