From 41454dcb3bcc00a5063b9eeb70a2344c3bf2c4fd Mon Sep 17 00:00:00 2001 From: zigzagcai Date: Tue, 16 Jul 2024 23:00:04 +0800 Subject: [PATCH] bug fix for test of state update with heads --- tests/ops/triton/test_selective_state_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/ops/triton/test_selective_state_update.py b/tests/ops/triton/test_selective_state_update.py index 696b2c77..e81807ae 100644 --- 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