I assume you mean using the module normally during training and then call torch.jit.trace at export time to create the JIT version. While this works for fully traceable models, for my use case I have mixed tracing and TorchScript, and I don’t see any way to disable the TorchScript annotations, the JIT is used for script functions/methods even if torch.jit.trace was never called. Is there any other way to achieve this?