I assume you’ve written a custom resnet model, as you are expecting two outputs?
Note that your current approach will not switch between the linear layers, but call new_logits on top of resnet.logits. Also, y_pred1 isn’t defined in your code.
If you haven’t defined a custom resnet but are using the torchvision implementation, note that this model does not contain a resnet.logits layer and you are simply assigning a new linear layer to this attribute.