import torchvision
import torch
from torch import nn



# train_data = torchvision.datasets.ImageNet("./data_image_net", split="train", download=True,
#                                            transform=torchvision.transforms.ToTensor)
from torch.utils.data import DataLoader

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)

print(vgg16_false)
print("——————————————————————")
print(vgg16_true)

dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)
vgg16_true.add_module("add_linear", nn.Linear(1000, 10))

vgg16_false.classifier[6] = nn.Linear(4096, 10)