I have these methods:
def set_jit_enabled(enabled: bool):
    """ Enables/disables JIT """
    if torch.__version__ < "1.7":
        torch.jit._enabled = enabled
    else:
        if enabled:
            torch.jit._state.enable()
        else:
            torch.jit._state.disable()


def jit_enabled():
    """ Returns whether JIT is enabled """
    if torch.__version__ < "1.7":
        return torch.jit._enabled
    else:
        return torch.jit._state._enabled.enabled

When training, I call set_jit_enabled(false) before instantiating modules that have JIT annotations. I haven’t found the proper way of doing this, but this hack works just fine.