Custom forward pass, but just for one epoch

I am looking for a way to run custom forward pass in an arbitrary model in Linear and Conv layers but for first epoch only.
Any ideas? modules() only work for a single sequence (doesn’t work if I have extra ops in model’s forward pass or multiple sequences running in parallel). Couldn’t find any way to use hooks (maybe missed something?) and the only idea I have right now is to unwrap grad_fn backwards to build forward sequence but it seems way too complicated.

I gues you should provide more information for me to help you. But it seems like you could solve your problem by implementing a custom forward pass within your model module, and then, use it given a condition within your forward pass.

Thank u!
Given is arbitrary model made of Linear and Conv layers.
During first epoch I want to compute new set of weights in each layer without backprop. Single forward pass is enough for this task (each layer requires values computed at layer before).

Problem is, list of modules does not provide correct order of execution (for example Sequential layer will appear in list of modules altogether with all layers from inside of it). And I’m not sure if it is possible to use forward hooks to modify weights.

Why don’t you just apply the forward hook to just Linear and Conv objectss?

accepted_modules = ['Linear', 'Conv']
for name, module in net.named_modules():
  if(name is in accepted_modules):
    module.register_forward_pre_hook(forward_pre_hook_fn)