From 203033ea16a288aa764c6d73bc3d9d9da6e6f87c Mon Sep 17 00:00:00 2001
From: duanjunwen <935724073@qq.com>
Date: Fri, 23 Aug 2024 08:57:27 +0000
Subject: [PATCH] [fix] fix weight not close;
---
tests/test_pipeline/test_schedule/test_zerobubble_pp.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/tests/test_pipeline/test_schedule/test_zerobubble_pp.py b/tests/test_pipeline/test_schedule/test_zerobubble_pp.py
index bf1fba3c6..b0927c0c4 100644
--- a/tests/test_pipeline/test_schedule/test_zerobubble_pp.py
+++ b/tests/test_pipeline/test_schedule/test_zerobubble_pp.py
@@ -389,7 +389,8 @@ def test_zerobubble_pipeline_base(
##########################
# fwd & bwd
output_base = model_base(input_base)
- loss_base = output_base.mean()
+ # loss_base = output_base.mean()
+ loss_base = criterion(output_base)
loss_base.backward()
print(f"After base fwd & bwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")