The problem is that your BN layers differ.
I used the following code to solve the problem (just override the train function of your model):
    def train(self, mode=True, freeze_bn=False, freeze_bn_affine=False):
        
        super(MyModel, self).train(mode)
        if freeze_bn:
            for m in self.modules():
                if isinstance(m, nn.BatchNorm2d):
                    '''Freezing Mean/Var of BatchNorm2D'''
                    m.eval()
                    if freeze_bn_affine:
                        '''Freezing Weight/Bias of BatchNorm2D'''
                        m.weight.requires_grad = False
                        m.bias.requires_grad = False