[GPU] Add an attribute to the torchscript model exported by metal (#4…

…7174) Summary: Pull Request resolved: #47174 As title ghstack-source-id: 115747991 Test Plan: Sandcastle Reviewed By: kimishpatel Differential Revision: D24616430 fbshipit-source-id: 2ccd264688471788f0dfea8bdc234fa69d39817f
pytorch · Nov 3, 2020 · ad3a3bd · ad3a3bd
1 parent 0ead9d5
commit ad3a3bd
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torch/csrc/jit/passes/metal_rewrite.cpp b/torch/csrc/jit/passes/metal_rewrite.cpp
@@ -190,6 +190,8 @@ script::Module metalOptimizeForMobile(
   metalFoldPrePackingOps(cloned_module);
   metalInsertCopyOps(cloned_module);
   removeDropout(cloned_module);
+  cloned_module.register_attribute(
+      "optimized_for_metal", BoolType::get(), true);
   return cloned_module;
 }