add #TODO

fattorib · Nov 10, 2024 · 1f9990e · 1f9990e
1 parent b0d8d98
commit 1f9990e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mamba/mamba_inner_fn.py b/mamba/mamba_inner_fn.py
@@ -162,7 +162,7 @@ def backward(ctx, grad_output):  # type: ignore
             x.mT.contiguous(), conv1d_w, conv1d_b, x_grad.contiguous(), act=1
         )
 
-        dx = torch.cat([dx_pre_conv.mT.contiguous(), dres], dim=-1)
+        dx = torch.cat([dx_pre_conv.mT.contiguous(), dres], dim=-1) #TODO: This is pretty slow, probably better to fuse with bwd on cc
 
         return (
             dx,