# First, we extend the VGG16 model with our new classifier
def E2EModel():
    e2e_model = models.vgg16(pretrained=True)
    e2e_model.train()

    classifier = list(e2e_model.classifier.children())[:-1]
    classifier.append(torch.nn.Linear(4096, 512))
    classifier.append(torch.nn.Sigmoid())
    classifier.append(torch.nn.Linear(512, 80))
    # Second sigmoid layer is dropped as it's included in the loss calculation
    e2e_model.classifier = nn.Sequential(*classifier)

    # Initialize final two linear layers with learned weights from Q5
    trained_tl_model = TwoLayerNet(4096, 512, 80)
    trained_tl_model.load_state_dict(torch.load('./outputs/q5.pth'))
    e2e_model.classifier[6].weight = trained_tl_model.linear1.weight
    e2e_model.classifier[8].weight = trained_tl_model.linear2.weight
    return e2e_model