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

…9131)
apache · Sep 27, 2021 · db8864b · db8864b
1 parent 0564d38
commit db8864b
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);