fix import

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

Author: pszemraj

samba_pytorch/modules/fused_rotary_embedding.py

@@ -4,10 +4,11 @@
 # Copyright (c) 2023, Tri Dao.
 
 
-import rotary_emb
 import torch
 from einops import rearrange
 
+import samba_pytorch.modules.rotary as rotary_emb
+
 
 class ApplyRotaryEmb(torch.autograd.Function):
     @staticmethod