careful with einsum non matmul cases

sonos · Aug 23, 2024 · 0d16734 · 0d16734
1 parent feb5aff
commit 0d16734
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/core/src/ops/einsum/mod.rs b/core/src/ops/einsum/mod.rs
@@ -268,7 +268,7 @@ impl TypedOp for EinSum {
         _outputs: &[&TypedFact],
     ) -> TractResult<AxesMapping> {
         let mut axes = self.axes.clone();
-        for input in [0, 1] {
+        for input in 0..inputs.len().min(2) {
             if inputs[input].datum_type.is_opaque() {
                 while axes.axes(InOut::In(input)).next().is_some() {
                     axes = axes.remove_axis_occurency(InOut::In(input), 0)?;