If you trace the model, you can just use torch.jit.is_tracing() to check whether you want the JIT-compatible bit and use the function otherwise.
There also is torch.jit.is_scripting() for the scripting equivalent.
Best regards
Thomas