Save the model as torchscript

[NimaDL]

Hi, I am trying to save the model as torchscript using the following line

model_scripted = torch.jit.script(net)

But when I added this line to snapshot function in logging script it raised an error

`def snapshot(self, net: torch.nn.Module, opt: Optimizer, sched: _LRScheduler = None, epoch: int = None, subdir='.')

    outfile = pt.join(self.dir, subdir, 'snapshot.pt')
    outfile_script = pt.join(self.dir, subdir, 'model_scripted.pt')
    if not pt.exists(os.path.dirname(outfile)):
        os.makedirs(os.path.dirname(outfile))

    model_scripted = torch.jit.script(net) # Export to TorchScript
    model_scripted.save(outfile_script) # Save

    torch.save(
        {'net': net.state_dict(), 'opt': opt.state_dict(), 'sched': sched.state_dict(), 'epoch': epoch}
        , outfile
    )
    return outfile`

Any idea how to save torchscript model inside the snapshot fucntion?

[NimaDL]

The problem solved!