github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-21 13:11:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

rename test

Browse Source
This commit is contained in:
oahzxl 2023-01-09 14:20:54 +08:00
parent cb68ee864a
commit 18a51c87fe
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_autochunk/test_autochunk.py → tests/test_autochunk/test_autochunk_codegen.py
View File

@ -18,9 +18,7 @@ def _test_fwd(model: torch.nn.Module, gm: ColoGraphModule, node, pair):
torch.cuda.reset_peak_memory_stats()
now_mem = torch.cuda.memory_allocated() / 1024**2
with torch.no_grad():
node1 = node.clone()
pair1 = pair.clone()
gm(node1, pair1)
gm(node.clone(), pair.clone())
new_now_mem = torch.cuda.memory_allocated() / 1024**2
new_max_mem = torch.cuda.max_memory_allocated() / 1024**2
print(