From 18a51c87fe0aa3a1210d7484fc09c16714e04bb7 Mon Sep 17 00:00:00 2001 From: oahzxl Date: Mon, 9 Jan 2023 14:20:54 +0800 Subject: [PATCH] rename test --- .../{test_autochunk.py => test_autochunk_codegen.py} | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) rename tests/test_autochunk/{test_autochunk.py => test_autochunk_codegen.py} (97%) diff --git a/tests/test_autochunk/test_autochunk.py b/tests/test_autochunk/test_autochunk_codegen.py similarity index 97% rename from tests/test_autochunk/test_autochunk.py rename to tests/test_autochunk/test_autochunk_codegen.py index 85a162084..1c5dd939d 100644 --- a/tests/test_autochunk/test_autochunk.py +++ b/tests/test_autochunk/test_autochunk_codegen.py @@ -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(