keep net type info when generating model complete net (pytorch#11032)

Summary:
keep net type info when generating model complete net. This will keep the performance optimization option
Pull Request resolved: pytorch#11032

Reviewed By: wat3rBro

Differential Revision: D9564125

Pulled By: harouwu

fbshipit-source-id: c6546af9b1d4ff5eddf6124e24a5da1b8baf47df

Author: harouwu

caffe2/python/model_helper.py

@@ -465,6 +465,8 @@ def GetCompleteNet(self):
         for op in new_net.Proto().op:
             op.debug_info = op.debug_info + "/param_init_net"
         new_net.AppendNet(self.net)
+        # keep the execution optimization
+        new_net.Proto().type = self.net.Proto().type
         return new_net
 
     def ConstructInitTrainNetfromNet(self, net):

caffe2/python/model_helper_test.py

@@ -8,6 +8,26 @@
 
 
 class ModelHelperTest(unittest.TestCase):
+    def test_get_complete_net_type(self):
+        model = model_helper.ModelHelper("test_orig")
+        brew.conv(
+            model,
+            "input",
+            "conv",
+            dim_in=3,
+            dim_out=16,
+            weight_init=("MSRAFill", {}),
+            kernel=3,
+            stride=1,
+            pad=0,
+        )
+        model.net.Proto().type = "async_scheduling"
+        net = model.GetCompleteNet()
+        model2 = model_helper.ModelHelper("test_new")
+        model2.ConstructInitTrainNetfromNet(net)
+        self.assertTrue(model2.net.Proto().type, "async_scheduling")
+        self.assertTrue(model2.param_init_net.Proto().type, "async_scheduling")
+
     def test_get_complete_net(self):
         model = model_helper.ModelHelper("test_orig")
         conv = brew.conv(