Solved by NikolaVetnic in post #6 
              
              
                So this is a solution I found to be working (please correct me if I am wrong): 
class NetBase(nn.Module):

    def __init__(self, model, num_classes):
        super().__init__()

        self.features = nn.ModuleList(model.children())[:-1]
        self.fc = nn.Linear(512, num_classes)

    def forwa…