Fix missing memory_efficient params

aliyun · Dec 13, 2024 · c34dbdd · c34dbdd
1 parent f16eafb
commit c34dbdd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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