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

rename test

pull/2364/head
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() torch.cuda.reset_peak_memory_stats()
now_mem = torch.cuda.memory_allocated() / 1024**2 now_mem = torch.cuda.memory_allocated() / 1024**2
with torch.no_grad(): with torch.no_grad():
node1 = node.clone() gm(node.clone(), pair.clone())
pair1 = pair.clone()
gm(node1, pair1)
new_now_mem = torch.cuda.memory_allocated() / 1024**2 new_now_mem = torch.cuda.memory_allocated() / 1024**2
new_max_mem = torch.cuda.max_memory_allocated() / 1024**2 new_max_mem = torch.cuda.max_memory_allocated() / 1024**2
print( print(