Update __init__.py

EleutherAI · Nov 26, 2023 · 0bd0839 · 0bd0839
1 parent 6ff6180
commit 0bd0839
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/megatron/fused_kernels/__init__.py b/megatron/fused_kernels/__init__.py
@@ -33,6 +33,7 @@ def load_fused_kernels():
     try:
         import scaled_upper_triang_masked_softmax_cuda
         import scaled_masked_softmax_cuda
+        import fused_rotary_positional_embedding 
     except (ImportError, ModuleNotFoundError) as e:
         print("\n")
         print(e)