github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 12:01:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

adapt new fx

Browse Source
This commit is contained in:
oahzxl
2023-01-10 11:56:00 +08:00
parent e532679c95
commit 7ab2db206f
4 changed files with 12 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/autochunk/estimate_memory.py
View File

@@ -28,12 +28,7 @@ class EstimateMemory(object):
return x
def _get_output_node(self, n):
fwd_out = {
x.uuid: x
for x in n.meta["fwd_out"]
if isinstance(x, torch.Tensor) and hasattr(x, "uuid")
}
out_size = activation_size(fwd_out)
out_size = activation_size(n.meta["fwd_out"])
out_node = [n.name] if out_size > 0 else []
return out_size, out_node