(model, preserved_attrs=[])
| 39 | |
| 40 | |
| 41 | def get_optimized_script(model, preserved_attrs=[]): |
| 42 | script = torch.jit.script(model) |
| 43 | if preserved_attrs != []: |
| 44 | script = torch.jit.freeze(script, preserved_attrs=preserved_attrs) |
| 45 | else: |
| 46 | script = torch.jit.freeze(script) |
| 47 | script = torch.jit.optimize_for_inference(script) |
| 48 | return script |
| 49 | |
| 50 | |
| 51 | def main(): |