From 0eb426a15c91e3b520454ae2ed81cfc35d649fb8 Mon Sep 17 00:00:00 2001
From: Dave Liddell <dliddell@xilinx.com>
Date: Tue, 6 Feb 2024 19:15:47 -0800
Subject: [PATCH] WIP: fixes for crashes

---
 lib/Dialect/Torch/IR/TorchOps.cpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp
index 1ff3db3d9890..04dd34c13e9a 100644
--- a/lib/Dialect/Torch/IR/TorchOps.cpp
+++ b/lib/Dialect/Torch/IR/TorchOps.cpp
@@ -2843,11 +2843,16 @@ OpFoldResult AtenIndexSelectOp::fold(FoldAdaptor adaptor) {
   auto self = getSelf();
   auto index = getIndex();
   auto selfTy = dyn_cast_or_null<ValueTensorType>(self.getType());
-  assert(selfTy);
   auto indexTy = dyn_cast_or_null<ValueTensorType>(index.getType());
-  assert(indexTy);
   auto resultTy = dyn_cast_or_null<ValueTensorType>(getType());
-  assert(resultTy);
+  if (!selfTy || !indexTy || !resultTy)
+    return nullptr;
+
+  if (!selfTy.hasSizes() || !indexTy.hasSizes() || !resultTy.hasSizes())
+    return nullptr;
+
+  if (!selfTy.hasDtype() || !indexTy.hasDtype() || !resultTy.hasDtype())
+    return nullptr;
 
   auto selfSizes = selfTy.getSizes();
   auto indexSizes = indexTy.getSizes();