make injective ops's opt schedule applied to every output tensor (#11820

)
apache · Jun 27, 2022 · 1115fd9 · 1115fd9
1 parent 6d6e070
commit 1115fd9
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/python/tvm/topi/x86/injective.py b/python/tvm/topi/x86/injective.py
@@ -74,12 +74,11 @@ def schedule_injective(outs):
         The computation schedule for the op.
     """
     outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
-    x = outs[0]
     s = te.create_schedule([x.op for x in outs])
     te.schedule.AutoInlineInjective(s)
-
-    if not is_empty_shape(x.shape):
-        schedule_injective_from_existing(s, x)
+    for x in outs:
+        if not is_empty_shape(x.shape):
+            schedule_injective_from_existing(s, x)
     return s