feat(aten::dropout_): Remove inplace dropout

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

Diff for: core/conversion/conversion_blacklist.cpp

@@ -22,7 +22,8 @@ const std::unordered_set<std::string>& get_non_convertable_nodes() {
         "prim::GetAttr",
         "prim::CallMethod",
         "prim::Drop",
-        "aten:dropout",
+        "aten::dropout",
+        "aten::dropout_"
     };
     return nonconvertable_nodes;
 }

Diff for: core/lowering/passes/remove_dropout.cpp

@@ -20,6 +20,20 @@ void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph) {
     remove_dropout.RegisterRewritePattern(
         dropout_pattern, no_dropout_pattern);
     remove_dropout.runOnGraph(graph);
+
+    std::string dropout_inplace_pattern = R"IR(
+        graph(%input, %4, %5):
+            %6 = aten::dropout_(%input, %4, %5)
+            return (%6))IR";
+    std::string no_dropout_inplace_pattern = R"IR(
+        graph(%input, %4, %5):
+            return (%input))IR";
+
+    torch::jit::SubgraphRewriter remove_dropout_inplace_pattern;
+    remove_dropout_inplace_pattern.RegisterRewritePattern(
+        dropout_inplace_pattern, no_dropout_inplace_pattern);
+    remove_dropout_inplace_pattern.runOnGraph(graph);
+
     LOG_GRAPH("Post remove dropout: " << *graph);
 }