From bd504c3cee0979450af4ce6db17c920aa5e96f88 Mon Sep 17 00:00:00 2001 From: xiaohou Date: Tue, 7 Dec 2021 11:17:47 +0800 Subject: [PATCH] [bin] check torch script (#36) * update run.sh * [network] add code to check whether the model can be exported to torch script or not --- kws/bin/train.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/kws/bin/train.py b/kws/bin/train.py index 8053b49..9e6117c 100644 --- a/kws/bin/train.py +++ b/kws/bin/train.py @@ -164,9 +164,9 @@ def main(): # !!!IMPORTANT!!! # Try to export the model by script, if fails, we should refine # the code to satisfy the script export requirements - # if args.rank == 0: - # script_model = torch.jit.script(model) - # script_model.save(os.path.join(args.model_dir, 'init.zip')) + if args.rank == 0: + script_model = torch.jit.script(model) + script_model.save(os.path.join(args.model_dir, 'init.zip')) executor = Executor() # If specify checkpoint, load some info from checkpoint if args.checkpoint is not None: