alleviate softmax datatype check (#24)

pai-disc · Nov 25, 2022 · 997551d · 997551d
1 parent 862576b
commit 997551d
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -321,9 +321,11 @@ class DecomposeAtenSoftmaxIntOp : public OpRewritePattern<AtenSoftmaxIntOp> {
   LogicalResult matchAndRewrite(AtenSoftmaxIntOp op,
                                 PatternRewriter &rewriter) const override {
     Value self = op.self();
-    if (!op.dtype().getType().isa<Torch::NoneType>())
-      return rewriter.notifyMatchFailure(
-          op, "Unimplemented non-None dtype for softmax");
+
+    // Do not need check dtype args here, since dtype have been infered in op.getType()
+    // if (!op.dtype().getType().isa<Torch::NoneType>())
+    //   return rewriter.notifyMatchFailure(
+    //       op, "Unimplemented non-None dtype for softmax");
 
     BaseTensorType tensorType = self.getType().cast<BaseTensorType>();
     if (!tensorType.hasDtype() || !tensorType.getDtype().isa<mlir::FloatType>())