Masaki Yano commited on
Commit
71c1fa1
·
1 Parent(s): 583f89f

chore(tools): add decode_in_inference option to export_torchscript (#1442)

Browse files
Files changed (1) hide show
  1. tools/export_torchscript.py +6 -1
tools/export_torchscript.py CHANGED
@@ -27,6 +27,11 @@ def make_parser():
27
  parser.add_argument("-expn", "--experiment-name", type=str, default=None)
28
  parser.add_argument("-n", "--name", type=str, default=None, help="model name")
29
  parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
 
 
 
 
 
30
  parser.add_argument(
31
  "opts",
32
  help="Modify config options using the command-line",
@@ -61,7 +66,7 @@ def main():
61
  if "model" in ckpt:
62
  ckpt = ckpt["model"]
63
  model.load_state_dict(ckpt)
64
- model.head.decode_in_inference = False
65
 
66
  logger.info("loading checkpoint done.")
67
  dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
 
27
  parser.add_argument("-expn", "--experiment-name", type=str, default=None)
28
  parser.add_argument("-n", "--name", type=str, default=None, help="model name")
29
  parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
30
+ parser.add_argument(
31
+ "--decode_in_inference",
32
+ action="store_true",
33
+ help="decode in inference or not"
34
+ )
35
  parser.add_argument(
36
  "opts",
37
  help="Modify config options using the command-line",
 
66
  if "model" in ckpt:
67
  ckpt = ckpt["model"]
68
  model.load_state_dict(ckpt)
69
+ model.head.decode_in_inference = args.decode_in_inference
70
 
71
  logger.info("loading checkpoint done.")
72
  dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])