I want a custom optimizer to react to the last layers of a graph in a specific way. Usually, the last group in a param_group is the last, but many networks have multiple output layers. Is there any way to know at the optimization step of the optimizer which ones are output layers?
All optimizers accept param_groups as argument to construct them.
So you can require the user to provide parameters with two groups called body and ouput in your optimizer init.
That way you know exactly which is which.