github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-25 15:01:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

rename

Browse Source
This commit is contained in:
oahzxl 2023-01-09 17:42:25 +08:00
parent a4ed5b0d0d
commit 865f2e0196
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/autochunk/search_chunk.py
View File

@ -48,7 +48,7 @@ class SearchChunk(object):
self.gm = gm self.gm = gm
self.print_mem = print_mem self.print_mem = print_mem
self.trace_indice = TraceIndice(list(gm.graph.nodes)) self.trace_indice = TraceIndice(list(gm.graph.nodes))
self.trace_indice.trace_index() self.trace_indice.trace_indice()
self.trace_flow = TraceFlow(self.trace_indice) self.trace_flow = TraceFlow(self.trace_indice)
self.reorder_graph = ReorderGraph(self.trace_indice) self.reorder_graph = ReorderGraph(self.trace_indice)
self.estimate_memory = EstimateMemory() self.estimate_memory = EstimateMemory()

2
colossalai/autochunk/trace_indice.py
View File

@ -525,7 +525,7 @@ class TraceIndice(object):
merge_to if i == merge_from else i for i in trace["indice"] merge_to if i == merge_from else i for i in trace["indice"]
] ]
def trace_index(self): def trace_indice(self):
for idx, node in enumerate(self.node_list): for idx, node in enumerate(self.node_list):
if node.op == "placeholder": if node.op == "placeholder":
self._assign_all_indice(node, idx) self._assign_all_indice(node, idx)