fix a bug which miss to patch rmsnorm

DeepLink-org · Nov 19, 2024 · 998040b · 998040b
1 parent f22692b
commit 998040b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/deeplink_ext/ops/rms_norm/__init__.py b/deeplink_ext/ops/rms_norm/__init__.py
@@ -11,3 +11,4 @@
 from .rms_norm_utils import import_RMSNorm, patch_RMSNorm
 
 MixedFusedRMSNorm = import_RMSNorm()
+patch_RMSNorm(MixedFusedRMSNorm)