bug fix for test of state update with heads (state-spaces#472)

mzusman · Jul 16, 2024 · c0a00bd · c0a00bd
1 parent 014c094
commit c0a00bd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/ops/triton/test_selective_state_update.py b/tests/ops/triton/test_selective_state_update.py
@@ -6,7 +6,7 @@
 import torch.nn.functional as F
 import pytest
 
-from einops import rearrange
+from einops import rearrange, repeat
 
 from mamba_ssm.ops.triton.selective_state_update import selective_state_update, selective_state_update_ref