bug fix

microsoft · Aug 7, 2024 · 7a16e5a · 7a16e5a
1 parent e7c805a
commit 7a16e5a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bitblas/ops/general_matmul/__init__.py b/bitblas/ops/general_matmul/__init__.py
@@ -128,7 +128,7 @@ def __initialize_propagate(self, propagate_a: Optional[TransformKind],
                 self.optimize_stratety == OptimizeStrategy.ContigousBatching
                 # TODO(lei): Should add ladder stage 3 inverse layout propagation in the expr.
                 # And recover the layout in the schedule template.
-                and self.M != 1 and 1 not in self.M):
+                and (self.M != 1 or (isinstance(self.M, Tuple) and 1 not in self.M))):
             object.__setattr__(self, "propagate_b", TransformKind.LDMatrixTransform)
 
         # TODO(lei): This is a limitation arose by pytorch and llvm