Torch jit script

I have fine-tuned a Huggingface (HF) ASR (xlsr_1b model). How do I convert to a jit.script model?
My finetuned HF model is saved as a .pt file as below:

model_id = “models/xlsr1b_torch_kas”
model = Wav2Vec2ForCTC.from_pretrained(model_id)
model_pt=“models/xlsr1b_torch_kas.pt”
chkpt={“model”: model.state_dict(),}
torch.save(chkpt, model_pt)

Now I want to convert this .pt model to jit. I load the saved .pt model first.

model_id = “models/xlsr1b_torch_kas”
model = Wav2Vec2ForCTC.from_pretrained(model_id)
checkpoint = torch.load(model_pt, map_location=“cpu”)
model.load_state_dict(checkpoint[“model”], strict=False)
model.eval()
model.to(“cpu”)
model.eval()

Now I want to convert to jit:

model.class.forward = torch.jit.ignore(model.class.forward)
model = torch.jit.script(model)

This is throwing error-

UnsupportedNodeError: function definitions aren’t supported:
File “/root/.virtualenvs/torchenv/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py”, line 455
if self._requires_grad and self.gradient_checkpointing and self.training:

            def create_custom_forward(module):
            ~~~ <--- HERE
            
                def custom_forward(*inputs):

I have tried using @torch.jit.ignore(drop=True) and @torch.jit.unused before
def create_custom_forward(module) but has not helped.

Where is the issue here as I think I am trying to ignore that part which is not required in the compiled version for inference.

Thanks for any pointer to resolution here.