From 9d516fa68f4e029d63b53d78803667bfa71e86d6 Mon Sep 17 00:00:00 2001
From: oahzxl <xuanlei.zhao@gmail.com>
Date: Sun, 18 Dec 2022 20:37:55 +0800
Subject: [PATCH] fix layernorm

---
 chunk_codegen.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/chunk_codegen.py b/chunk_codegen.py
index 5e2130ee7..77c28fd32 100644
--- a/chunk_codegen.py
+++ b/chunk_codegen.py
@@ -574,7 +574,7 @@ class IndexTracer(object):
             node_idx (int)
         """
         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):
         """