You can also include other parameters in your forward function such as a boolean value whether or not to use head1 or head2.
def forward(self, input, use_head1=True):
x = sharedNet(input)
if useHead1:
   x = head1(x)
else:
   x = head2(x)

output = model(input, useHead1=False) #Use head2