mirror of https://github.com/hpcaitech/ColossalAI
fix layernorm
parent
e66a18a0bf
commit
9d516fa68f
|
@ -574,7 +574,7 @@ class IndexTracer(object):
|
||||||
node_idx (int)
|
node_idx (int)
|
||||||
"""
|
"""
|
||||||
self._assign_index_as_input(node, idx)
|
self._assign_index_as_input(node, idx)
|
||||||
self._mark_computation(node, idx, [-1, -2])
|
self._mark_computation(node, idx, [-1])
|
||||||
|
|
||||||
def _assign_elementwise_index(self, node, idx):
|
def _assign_elementwise_index(self, node, idx):
|
||||||
"""
|
"""
|
||||||
|
|
Loading…
Reference in New Issue