Unable to save the pytorch model using torch.jit.scipt

jit
1

Hi Team,

I’m trying to save the model after every epoch. Unfortunately, I’m unable to save and I’m getting the below error.

This is the git repo I’m using GitHub - CSAILVision/semantic-segmentation-pytorch: Pytorch implementation for Semantic Segmentation/Scene Parsing on MIT ADE20K dataset

in their train script, they are trying to save the model like below,

(net_encoder, net_decoder, crit) = nets

dict_encoder = net_encoder.state_dict()
dict_decoder = net_decoder.state_dict()

  torch.save(
      dict_encoder,
      '{}/encoder_epoch_{}.pth'.format(cfg.DIR, epoch))
  torch.save(
      dict_decoder,
      '{}/decoder_epoch_{}.pth'.format(cfg.DIR, epoch))

however, I have changed those bits to torch script manner like below,

scripted_encoder = torch.jit.script(net_encoder)
scripted_encoder.save('encoder_scripted.pt') 

scripted_decoder = torch.jit.script(net_decoder)
scripted_decoder.save('decoder_scripted.pt')

when I try to save the model as torch script way, I get the following error.

Screenshot from 2022-06-17 00-56-52
Screenshot from 2022-06-17 00-56-521920×1080 333 KB

@ptrblck @ppwwyyxx could you please help me to solve this issue and save the model using torch script.

Sytem details:

 - PyTorch Version: 1.4.0
 - Torchvision Version: 0.5.0
 - OS: Ubuntu 18.04.3 LTS
 - How you installed PyTorch: `pip`
 - Python version: 3.6.13
 - CUDA/cuDNN version: CUDA Version: 11.0
 - GPU models and configuration: NVIDIA K80