Specify weights when loading the model in the pkl to pb conversion tool

Reviewed By: rbgirshick Differential Revision: D7484464 fbshipit-source-id: f92804fedc3ae040e585e7c0c2eb3e44526eef22
ronghanghu · Apr 3, 2018 · 1d67750 · 1d67750
1 parent 22636d4
commit 1d67750
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tools/convert_pkl_to_pb.py b/tools/convert_pkl_to_pb.py
@@ -333,7 +333,7 @@ def _save_models(all_net, all_init_net, args):
 
 
 def load_model(args):
-    model = test_engine.initialize_model_from_cfg()
+    model = test_engine.initialize_model_from_cfg(cfg.TEST.WEIGHTS)
     blobs = mutils.get_ws_blobs()
 
     return model, blobs