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

…pache#9131)
ylc · Sep 29, 2021 · e019e76 · e019e76
1 parent 93cdf4a
commit e019e76
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);