8192 upper

Author: mayank31398

megatron/fused_kernels/tests/test_fused_kernels.py

@@ -301,7 +301,7 @@ class DummyArgs:
         print("\n[Fail] Please install `transformers` package to test fused kernels\n")
         exit(-1)
 
-    from megatron.fused_kernels.cuda import load
+    from megatron.fused_kernels import load
     load(DummyArgs())
 
     test_load_fused_kernels()