I have a pretrained model that takes an input of vision and other state data and splits it. The vision data passes through an encoder before returning to the rest of the model, but this requires a fair amount of processing power and will be on relatively slow hardware. I have access to a separate processor that I want to offload the encoding task to, but for this, I would need to separate the model so that the encoder only gets the vision input, and the output is then fed to the rest of the model alongside the other state data inputs.
I haven’t found many sources that are helpful for this specific case and was hoping for some advice. I’ve already separated out the encoder by creating a new module for it, but my issue is saving the rest of the model in a way that it accepts the encoder output as an input.
Thank you in advance.