Freezing only part of the model

saar-eliad · June 24, 2021, 7:05am

I want to change the following

mlcommons/inference/blob/master/speech_recognition/rnnt/pytorch_SUT.py#L73-L83

    
      
          model.eval()
          model.encoder = torch.jit.script(model.encoder)
          model.encoder = torch.jit._recursive.wrap_cpp_module(
              torch._C._freeze_module(model.encoder._c))
          model.prediction = torch.jit.script(model.prediction)
          model.prediction = torch.jit._recursive.wrap_cpp_module(
              torch._C._freeze_module(model.prediction._c))
          model.joint = torch.jit.script(model.joint)
          model.joint = torch.jit._recursive.wrap_cpp_module(
              torch._C._freeze_module(model.joint._c))
          model = torch.jit.script(model)

(RNN-T from mlcommons benchmark)

to use the new torch.jit.freeze API, but it throws errors.
Can this be done?