[Torch] add constriants when decompose aten.split_with_sizes (llvm#3555)

vinayakdsci · Jul 23, 2024 · 7884642 · 7884642
1 parent 45c85c3
commit 7884642
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -1352,7 +1352,7 @@ class DecomposeAtenSplitWithSizesOp
 
     auto sliceTy =
         dyn_cast_or_null<Torch::BaseTensorType>(resultTy.getContainedType());
-    if (!isa<Torch::BaseTensorType>(sliceTy))
+    if (!sliceTy || !sliceTy.hasSizes())
       return rewriter.notifyMatchFailure(op, "Slice type is unknown");
 
     int64_t dimInt = 0;