Best practice for adding an extra layer after QAT

[Anurag Ranjan (anuragranj)]

I have a model that undergoes a QAT using AIMET.

  quant_sim = QuantizationSimModel(
      model,
      dummy_input=dummy_input,
      quant_scheme=QuantScheme.training_range_learning_with_tf_enhanced_init,
      default_param_bw=8,
      default_output_bw=8
  )

  quant_sim.compute_encodings(forward_loop, forward_pass_callback_args=model)

# ... QAT training in pytorch

Once the model is trained, i would like to add a preprocessing layer to the model.

new_model = torch.nn.Sequential(extra_module, model)

How should I use the new_model with quant_sim so I don't have to run the QAT again?