Hi @Caruso,
The simplest way to explain what I want to do is replace the module of the pre-trained model with a different one such that the pre-trained model still works.
For example referencing what I was trying:

Replace the last fully connected net with a fully connected net with different hyper parameters (e.g. 5 labels rather than 10000). [succeeded with model.fc = torch.nn.Linear(in_features=512, out_features=fc_out_features, bias=True)]
Replace every batch norm with a batch norm that does not track running means but instead uses mini-batch statistics. [Failed, idk how to do yet]

For my specific case a copy of the same module with different hyperparams would do (but ideally replacing any module with a different one is what would be more general and robust).
Thanks for taking the time to read my attempt. I will go ahead and read the links I pasted since those seem to be very related to my challenge.