Add type checking for lhs and rhs

Author: erman-gurses

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -825,14 +825,16 @@ void AtenAddTensorOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
     if (!lhs.hasValue() || !rhs.hasValue())
       return failure();
 
+    if (!lhs.getValue().getType().isa<Torch::IntType>() ||
+        !rhs.getValue().getType().isa<Torch::IntType>())
+      return failure();
+
     Value mul =
         rewriter.create<AtenMulIntOp>(op->getLoc(), rhs.getValue(), op.alpha());
-    if (!mul.getType().isa<Torch::IntType>())
-      return failure();
+
     Value add =
         rewriter.create<AtenAddIntOp>(op->getLoc(), lhs.getValue(), mul);
-    if (!add.getType().isa<Torch::IntType>())
-      return failure();
+
     rewriter.replaceOpWithNewOp<PrimNumToTensorScalarOp>(
         op, op.self().getType(), add);
     return success();