github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[fix] fix weight not close;

pull/6034/head
duanjunwen 2024-08-23 08:57:27 +00:00
parent c18ef060cf
commit 203033ea16
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_pipeline/test_schedule/test_zerobubble_pp.py
View File

@ -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;")