Got stuck in TorchScript Trace

jit
2

You would have to torch.jit.script the model instead of tracing.
The next stable release will use the nvfuser backend during scripting which has dynamic shape support.