Error using torch.save after loaded and traced hubert model with infer.lib.jit.to_jit_model

[simosechi]

I need to bring the model in c++ but I'm encountering some problem in saving hubert using TorchScript.
The following error appeare after I run torch.save for the traced model.

Could not export Python function call 'GradMultiply'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__

The error is reproducible using the following code

import os
import sys

now_dir = os.getcwd()
sys.path.append(now_dir)

from infer.lib.jit import to_jit_model

hb_model, traced = to_jit_model("assets/hubert/hubert_base.pt", "hubert", "trace", "assets/hubert/hubert_inputs.pth")

#after creating a folder traced
traced.save("traced/hubert_traced.pt")

I understand that maybe the problem is related with lib fairseq also, but I hope it can be resolved here