More if compaction

daveliddell · Feb 7, 2024 · 00d0817 · 00d0817
1 parent 3f7f475
commit 00d0817
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp
@@ -2854,11 +2854,8 @@ OpFoldResult AtenIndexSelectOp::fold(FoldAdaptor adaptor) {
   auto indexSizes = indexTy.getSizes();
   auto resultSizes = resultTy.getSizes();
 
-  if (selfTy.getDtype() != resultTy.getDtype())
-    return nullptr;
-  if (selfSizes.size() != resultSizes.size())
-    return nullptr;
-  if (indexSizes.size() != 1)
+  if (selfTy.getDtype() != resultTy.getDtype() ||
+      selfSizes.size() != resultSizes.size() || indexSizes.size() != 1)
     return nullptr;
 
   // If the selection results in a tensor of the same dimensions as the