From 32e642bf400237512fcc473c21b648565b7684e1 Mon Sep 17 00:00:00 2001 From: Wenhao Chen Date: Fri, 1 Mar 2024 16:50:58 +0800 Subject: [PATCH] revert: enable return_outputs when necessary --- tests/test_shardformer/test_model/_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_shardformer/test_model/_utils.py b/tests/test_shardformer/test_model/_utils.py index 85be9a242..f5fc21b4c 100644 --- a/tests/test_shardformer/test_model/_utils.py +++ b/tests/test_shardformer/test_model/_utils.py @@ -181,6 +181,7 @@ def run_forward_backward_with_hybrid_plugin( data_iter = iter([data]) sharded_output = booster.execute_pipeline( + data_iter, sharded_model, _criterion, sharded_optimizer, return_loss=True, return_outputs=True data_iter, sharded_model, _criterion,