class modelA(Module):
    def __init__(self):
        super().__init__()
        self.base= ...
        self.headA= ...
    
    def forward(self,input):
        x= self.base(input)
        outA= self.headA(x)

        return outA

class modelB(Module):
    def __init__(self):
        super().__init__()
        self.base=...
        self.headB= ...

    def forward(self,input):
        x= self.base(input)
        outB= self.headB(x)
        return outB

this model was supposed to work in way that weight of base layer should be shared, so as per suggestion here 2 by @apaszke I defined a combined model as follow
class model_combined(Module):
    def __init__(self):
        super().__init__()
        self.base=...
        self.headA= ...
        self.headB= ...

    def forward(self,input):
        x= self.base(input) # this need to be freezed
        outA= self.headA(x)
        y=self.base(input) #this neeeds to backpropagte
        outB= self.headB(y)
        return outA, outB

in training loop I am calculating loss as
total_loss= loss(outA)+ loss(outB)
total_loss.backward()

but after some epoch, I need to freeze modelA (i.e self.headA in model_combined)
My questions are:

since base is common layer here, so how to freeze headA (weight of headA should not change) such that base is updated due to backpropagation of loss(outB) but not due to loss(outA)

will it be a bad practice if instead of freezing I just don’t backpropagate loss through headA i.e change total_loss= loss(outB)