[quant][graphmode][fix] cloning schema in insert_observers (#40624)

jerryzh168 · facebook-github-bot · commit 0309f6a4bb2b · 2020-06-26T20:19:09.000-07:00
Summary: Pull Request resolved: #40624 Previously we didn't clone schema, so the default schema is used, this is causing issue for some models Test Plan: Imported from OSS Differential Revision: D22259519 fbshipit-source-id: e2a393a54cb18f55da0c7152a74ddc22079ac350
diff --git a/torch/csrc/jit/passes/quantization/insert_observers.cpp b/torch/csrc/jit/passes/quantization/insert_observers.cpp
@@ -242,12 +242,18 @@ class ModuleCloneHelper {
     remapTypes(graph.get(), source, target, module_qconfig_map, type_remap_fn);
     // remap self
     graph->inputs()[0]->setType(target.type());
+    // we only support %self being Module in the arguments of function
+    auto schema_type_remap_fn = [&](TypePtr type_ptr) {
+      return type_remap_fn(type_ptr, module_qconfig_map.at(source._ivalue()));
+    };
+    auto schema =
+        method.getSchema().cloneWithRemappedTypes(schema_type_remap_fn);
     const auto this_method_name =
         c10::QualifiedName(*target.type()->name(), method.name());
     auto copied = target._ivalue()->compilation_unit()->create_function(
         this_method_name, graph);
     target.type()->addMethod(copied);
-    // we'll use default schema for cloned method
+    copied->setSchema(std::move(schema));
   }
 };
 
