Hello @driazati
I work  with @Sebastien_Fischman and indeed it works in 1.4.0 now.
But we still have a problem : the models we trained beforehand on version 1.0.1 and saved as a JIT does not work anymore, when trying to load them.
I searched for a while, and I found that inside the JIT the problematic code was
 _12 = torch.index(_11, [annotate(Tensor, None), cont_idxs])
With new version, code would be something like
 _12 = torch.index(_11, annotate(List[Optional[Tensor]],[None, cont_idxs]))
If I changed manually the code for this new line, it works.
Problem is we are using PyTorch models in production, and manually modifying the models will be an issue.
Do you have a migration script to update the JIT file ? or any solution for JIT to be retrocompatible ?