diff --git a/workload_generator/mocked_model/AiobMegatron.py b/workload_generator/mocked_model/AiobMegatron.py
index 60e10d3..599107d 100755
--- a/workload_generator/mocked_model/AiobMegatron.py
+++ b/workload_generator/mocked_model/AiobMegatron.py
@@ -361,7 +361,7 @@ def __init__(self, args=None):
     @cuda_timing_decorator
     def _apply(self, hidden_states):
         output_lay = FastLayerNormFN.apply(
-            hidden_states, self.lay_weight, self.bias, 1e-05
+            hidden_states, self.lay_weight, self.bias, 1e-05, 1
         )
 
         return output_lay