github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-28 08:17:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

rename index tracer

Browse Source
This commit is contained in:
oahzxl 2022-12-06 17:35:27 +08:00
parent a9d64377bb
commit 6d99994a7a
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
chunk_codegen.py
View File

@ -17,7 +17,7 @@ def _delete_free_var_from_last_use(user_to_last_uses):
user_to_last_uses[key].remove(n) user_to_last_uses[key].remove(n)
class NodeIndexTracer(object): class IndexTracer(object):
def __init__(self, gm) -> None: def __init__(self, gm) -> None:
self.gm = gm self.gm = gm
self.nodes_list = list(gm.graph.nodes) self.nodes_list = list(gm.graph.nodes)
@ -683,7 +683,7 @@ class ChunkRegionSearch(object):
self.gm = gm self.gm = gm
self.node_list = list(gm.graph.nodes) self.node_list = list(gm.graph.nodes)
self.memory_estimator = MemoryEstimator() self.memory_estimator = MemoryEstimator()
self.index_tracer = NodeIndexTracer(gm) self.index_tracer = IndexTracer(gm)
self.index_tracer.trace_node_idx() self.index_tracer.trace_node_idx()
def _find_peak_node(self, mem_peak): def _find_peak_node(self, mem_peak):