revert: enable return_outputs when necessary

2024-03-01 16:50:58 +08:00 · 2024-03-01 16:50:58 +08:00 · 32e642bf40
parent 856b39f69d
commit 32e642bf40
1 changed files with 1 additions and 0 deletions
--- 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,