Dynamically change a models forward function during runtime

I am attempting to search for good branches in a model by dynamically adding a branch in the forward function by editing the forward function ast, unparsing the updated ast and writing it back to the file where the model is stored, and making a new object based on the updated model, with the old weights reloaded. But I can’t get the model to use the new forward function unless the process is completely shutdown and manually restarted. Is it even possible to edit a model in such a manner, or to reload its forward function at runtime?

An example workflow of how I’m attempting it is:

  1. Have a model ModelA in models.py
  2. In training.py, train the model
  3. In generation.py, take a reference to the model, get the ast of the model, and edit the forward function
  4. Unparse the edited ast and rewrite to models.py
  5. Save the weights of the model
  6. Reload models.py with importlb.reload
  7. Create a new model object
  8. Load the saved weights to the new object
  9. Train again to train the branch layers

After using importlib.reload, the model still has the old forward function in memory and does not use the updated one unless I kill the script and restart it, but I need to reload during runtime without killing the script.

Well I think you can do that without writting the file
if you call “eval” over a text string you should be able to create a working function.
then you just need to assign self.forward=new_fn and that’s all
You can ofc dump the function somewhere but that’s more logging that a requirement.

Another options is ofc given your new files a different name so you skip all the import issues.

Sorry for getting back so late, I was actually just being stupid. I was still recreating the object with the old non reloaded module. I would have rathered not to write the new model to a file, but I need to show how the model is being affected for what I’m doing.