AI/wekws
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bin] check torch script (#36)

Browse Source 
* update run.sh

* [network] add code to check whether the model can be exported to torch script or not
This commit is contained in:
xiaohou 2021-12-07 11:17:47 +08:00 committed by GitHub
parent afbc1d2960
commit bd504c3cee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
kws/bin/train.py
View File

@ -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: