github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-19 13:50:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fix] fix weight not close;

Browse Source
This commit is contained in:
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 # fwd & bwd
output_base = model_base(input_base) output_base = model_base(input_base)
loss_base = output_base.mean() # loss_base = output_base.mean()
loss_base = criterion(output_base)
loss_base.backward() loss_base.backward()
print(f"After base fwd & bwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;") print(f"After base fwd & bwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")