From 78846425e24479cff6e8499495656443e406efc0 Mon Sep 17 00:00:00 2001
From: Yuanqiang Liu <liuyuanqiang.yqliu@bytedance.com>
Date: Tue, 23 Jul 2024 10:34:29 +0800
Subject: [PATCH] [Torch] add constriants when decompose aten.split_with_sizes
 (#3555)

---
 lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
index 46b218535c67..804683ac2e1c 100644
--- 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;