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

rename index tracer

pull/2364/head
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)
class NodeIndexTracer(object):
class IndexTracer(object):
def __init__(self, gm) -> None:
self.gm = gm
self.nodes_list = list(gm.graph.nodes)
@ -683,7 +683,7 @@ class ChunkRegionSearch(object):
self.gm = gm
self.node_list = list(gm.graph.nodes)
self.memory_estimator = MemoryEstimator()
self.index_tracer = NodeIndexTracer(gm)
self.index_tracer = IndexTracer(gm)
self.index_tracer.trace_node_idx()
def _find_peak_node(self, mem_peak):