Update initialize.py

EleutherAI · Nov 26, 2023 · eed2204 · eed2204
1 parent 6e8e92a
commit eed2204
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/megatron/initialize.py b/megatron/initialize.py
@@ -57,6 +57,7 @@ def finish_mpu_init():
     if (
         neox_args.scaled_upper_triang_masked_softmax_fusion
         or neox_args.scaled_masked_softmax_fusion
+        or neox_args.rope_fusion
     ):
         fused_kernels.load_fused_kernels()