[DeltaNet] fix a bug of recurrent mode

fla-org · Oct 1, 2024 · aef77f3 · aef77f3
1 parent 376ac58
commit aef77f3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fla/ops/delta_rule/fused_recurrent.py b/fla/ops/delta_rule/fused_recurrent.py
@@ -266,7 +266,7 @@ def forward(ctx, q, k, v, beta, scale=None, initial_state=None, output_final_sta
             num_warps=num_warps,
             num_stages=num_stages,
         )
-        o = o.squeeze(1)
+        o = o.squeeze(0)
         ctx.save_for_backward(q, k, v, beta, initial_state)
         ctx.scale = scale
         return o, final_state