Fix bug in flatbuffer deserialization

Pull Request resolved: pytorch#78344 Approved by: https://github.com/qihqi
zmmwl · May 31, 2022 · 7e72c96 · 7e72c96
1 parent f42b42d
commit 7e72c96
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/test/cpp/jit/test_flatbuffer.cpp b/test/cpp/jit/test_flatbuffer.cpp
@@ -16,6 +16,7 @@
 #include <torch/csrc/jit/mobile/module.h>
 #include <torch/csrc/jit/mobile/parse_bytecode.h>
 #include <torch/csrc/jit/mobile/parse_operators.h>
+#include <torch/csrc/jit/operator_upgraders/upgraders.h>
 #include <torch/csrc/jit/serialization/export.h>
 #include <torch/csrc/jit/serialization/export_bytecode.h>
 #include <torch/csrc/jit/serialization/flatbuffer_serializer.h>
@@ -837,6 +838,7 @@ TEST(FlatbufferTest, DuplicateSetState) {
 
   auto buff = save_mobile_module_to_bytes(bc);
   mobile::Module bc2 = parse_mobile_module(buff.data(), buff.size());
+  ASSERT_TRUE(torch::jit::is_upgraders_map_populated());
   const auto methods2 = bc.get_methods();
   ASSERT_EQ(methods2.size(), expected_n);
 }

diff --git a/torch/csrc/jit/serialization/flatbuffer_serializer_jit.cpp b/torch/csrc/jit/serialization/flatbuffer_serializer_jit.cpp
@@ -1,6 +1,7 @@
 #include <torch/csrc/jit/serialization/flatbuffer_serializer_jit.h>
 
 #include <torch/csrc/jit/mobile/flatbuffer_loader.h>
+#include <torch/csrc/jit/operator_upgraders/upgraders_entry.h>
 #include <torch/csrc/jit/serialization/export.h>
 #include <torch/csrc/jit/serialization/export_bytecode.h>
 #include <torch/csrc/jit/serialization/flatbuffer_serializer.h>
@@ -14,6 +15,9 @@ Module parse_and_initialize_jit_module(
     size_t size,
     ExtraFilesMap& extra_files,
     c10::optional<at::Device> device) {
+#if ENABLE_UPGRADERS
+  populate_upgraders_graph_map();
+#endif
   auto* flatbuffer_module = mobile::serialization::GetMutableModule(data.get());
   FlatbufferLoader loader;
   mobile::Module mobilem = loader.parseModule(flatbuffer_module);