feat(aten::dropout_): Remove inplace dropout

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Jun 25, 2020 · 7aa57c3 · 7aa57c3
1 parent 19c91f2
commit 7aa57c3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/core/conversion/conversion_blacklist.cpp b/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 --git a/core/lowering/passes/remove_dropout.cpp b/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);
 }