commit b849f101b17ec56b165b39407406a61529b2058a
author James Reed <jamesreed@fb.com> Fri Jun 07 17:18:58 2019 -0700
committer Facebook Github Bot <facebook-github-bot@users.noreply.github.com> Fri Jun 07 17:28:55 2019 -0700
tree f2d76ae805f01a406eabd6208a887dd80e76717a
parent 66d596645a66ee503bdd3ea0a9ec6f51a2bd8399

Fix slow unpickling (#21542)

Summary:
This was looking at the number of elements in the memo table, not the total capacity, and was thus calling reserve() a lot more than it should have
Pull Request resolved: https://github.com/pytorch/pytorch/pull/21542

Reviewed By: driazati

Differential Revision: D15723132

Pulled By: jamesr66a

fbshipit-source-id: 20e1f9099b6a51a33994ea9dbc3f22eb3bc0c8f9

torch/csrc/jit/pickler.cpp

diff --git a/torch/csrc/jit/pickler.cpp b/torch/csrc/jit/pickler.cpp
index 9f6e2a4..5803a91 100644
--- a/torch/csrc/jit/pickler.cpp
+++ b/torch/csrc/jit/pickler.cpp
@@ -600,7 +600,7 @@
     } break;
     case OpCode::BINPUT: {
       size_t memo_id = read<uint8_t>();
-      if (memo_table_.size() <= memo_id) {
+      if (memo_table_.capacity() <= memo_id) {
         memo_table_.reserve(1 + 2 * memo_id);
       }
       memo_table_.push_back(stack_.back());
@@ -612,7 +612,7 @@
           "Found a LONG_BINPUT opcode, but size_t on this system is "
           "not big enough to decode it");
       size_t memo_id = read<uint32_t>();
-      if (memo_table_.size() <= memo_id) {
+      if (memo_table_.capacity() <= memo_id) {
         memo_table_.reserve(1 + 2 * memo_id);
       }
       memo_table_.push_back(stack_.back());