Update optimum/habana/transformers/models/mt5/modeling_mt5.py

Co-authored-by: Yaser Afshar <[email protected]>
Gaurav7888 · Aug 5, 2024 · 49a234e · 49a234e
1 parent 5aa891d
commit 49a234e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/optimum/habana/transformers/models/mt5/modeling_mt5.py b/optimum/habana/transformers/models/mt5/modeling_mt5.py
@@ -34,9 +34,11 @@
 
 try:
     from habana_frameworks.torch.hpex.normalization import FusedRMSNorm as FusedRMSNorm
+
+    has_fused_rms_norm = True
 except ImportError:
+    has_fused_rms_norm = False
     print("Not using HPU fused kernel for RMSNorm")
-    FusedRMSNorm = None
 
 def gaudi_mt5_layernorm_forward(self, hidden_states):
     """