I guess there is a better way to do it and would eventually be updated in the thread, but as far as I think adding a Linear layer in front of the input node of the pre-trained model should do the trick.
class ModelB(torch.nn.Module):
    def __init__(self):
        self.fc0 = torch.nn.Linear(new_shape,4)
        self.pretrained = ModelA()

    def forward(self, x):
         return self.pretrained(self.fc0(x))