unstack usage requires continuity

swfsql · Feb 9, 2024 · 1207867 · 1207867
1 parent 1005250
commit 1207867
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dfdx/src/nn/layers/mamba_minimal.rs b/dfdx/src/nn/layers/mamba_minimal.rs
@@ -673,7 +673,7 @@ pub mod stateless {
             Err(_delta_a) => unreachable!(),
         };
         let (delta_a, _delta_a_tape): (Vec<Tensor<(Batch, DInner, DState), _, _, _>>, _) =
-            delta_a.try_unstack()?;
+            delta_a.try_contiguous()?.try_unstack()?;
         //
         // delta B
         let delta_bu: Tensor<(usize, Batch, DInner, DState), _, _, _> = match delta_bu.try_realize()
@@ -682,7 +682,7 @@ pub mod stateless {
             Err(_delta_bu) => unreachable!(),
         };
         let (delta_bu, _delta_bu_tape): (Vec<Tensor<(Batch, DInner, DState), _, _, _>>, _) =
-            delta_bu.try_unstack()?;
+            delta_bu.try_contiguous()?.try_unstack()?;
         //
         // C
         let c: Tensor<(usize, Batch, DState, C1), _, _, _> = match c