tjoseph:

Can you explain why it does not make sense?


JIT speedups mostly come from tensor operations - like fused math operations and some optimizations specific to tensors, ops or layers. Classes like above are closer to non-computational wrappers - similar to training loop code, there is not much to optimize there.



 tjoseph:

Do you know whether torch.jit.script will track tensor attributes (i.e. requires_grad) or is that something that the jit is generally not supposed to do and it is only meant for deployment?


It is read-only (and used as meta-information in compilation). Instead, you can detach() tensors, and for parameters - you’d have to set it from outside.



 tjoseph:

If tracing would include tensor attributes I could just pass my whole forward pass to jit.trace which would properly separate optimization from model semantics. jit.script would be fine since it can be used at least with decorators, but unfortunately the missing support of torch.distributions does not allow jit.script.


Dunno, I find jit.trace too dangerous / limiting to use on module trees. As for torch.distributions, I mostly stopped using these wrappers when coding for performance.