Fix the missing dtype attribute of tir.Shuffle in Python level (a…

…pache#9131)
ylc · Jan 13, 2022 · ca87c28 · ca87c28
1 parent 8ac0d3c
commit ca87c28
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/include/tvm/tir/expr.h b/include/tvm/tir/expr.h
@@ -947,6 +947,7 @@ class ShuffleNode : public PrimExprNode {
   Array<PrimExpr> indices;
 
   void VisitAttrs(AttrVisitor* v) {
+    v->Visit("dtype", &dtype);
     v->Visit("vectors", &vectors);
     v->Visit("indices", &indices);
     v->Visit("span", &span);