commit 37c3d503b7d6b01435f44e4ecb26b6e579df5743
author Will Constable <whc@meta.com> Thu Aug 01 11:18:26 2024 -0700
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Fri Aug 02 20:38:17 2024 +0000
tree 81a01d70fdaf2ffbf6ae11ae7da5515d98b47e00
parent 7c1cca9fdaf216f303d064c9558f07803369b915

[pipelining] Make test_schedule quiet (#132369)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/132369
Approved by: https://github.com/H-Huang
ghstack dependencies: #129810, #130378

test/distributed/pipelining/test_schedule.py

diff --git a/test/distributed/pipelining/test_schedule.py b/test/distributed/pipelining/test_schedule.py
index 6114401..b2fed7f 100644
--- a/test/distributed/pipelining/test_schedule.py
+++ b/test/distributed/pipelining/test_schedule.py
@@ -105,7 +105,12 @@
                 if num_microbatches % group_size != 0:
                     continue
 
-                print(f"{num_local_stages=} {num_microbatches=} {group_size=}")
+                logger.info(
+                    "num_local_stages=%d num_microbatches=%d group_size=%d",
+                    num_local_stages,
+                    num_microbatches,
+                    group_size,
+                )
                 num_stages = num_local_stages * group_size
                 stages = [
                     MockPipelineStage(group_size=group_size, num_stages=num_stages)