diff --git a/llvm/projects/keras/frontend/hpvm_dfg_translator.py b/llvm/projects/keras/frontend/hpvm_dfg_translator.py
index e66ed64dff3aa5d0b0c512406c7b7e4e90a6eb31..066d5d23b90fae0048036b69bcdd7556adec9c92 100644
--- a/llvm/projects/keras/frontend/hpvm_dfg_translator.py
+++ b/llvm/projects/keras/frontend/hpvm_dfg_translator.py
@@ -365,7 +365,7 @@ class HPVMTranslator:
     inst_str += "  void *r = __visc__tensor_batchnorm(t1, t2, t3, t4, t5, "
     inst_str += str(cur_node.epsilon) + "); \n"
 
-    footer_str = self.genNodeFooter(5)
+    footer_str = self.genNodeFooter(2)
     inst_str += footer_str
         
     self.node_str += inst_str