VGG parameters not updating
I am trying to train the VGG network from PyTorch to build a predictive model for the FashionMNIST dataset. But when I print the gradients out, it seems that the parameters are not updating and the gradient is always zero. Here is my implementation
## Specify Batch Size
train_batch_size = 32
test_batch_size = 32

## Specify Image Transforms
img_transform = transforms.Compose([
    transforms.Resize((64,64)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

## Download Datasets
train_data = FashionMNIST('./data', transform=img_transform, download=True, train=True)
test_data = FashionMNIST('./data', transform=img_transform, download=True, train=False)

## Initialize Dataloaders
training_dataloader = DataLoader(train_data, batch_size=train_batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=test_batch_size, shuffle=True)

vgg16 = models.vgg16()
model_a = vgg16
model_a.classifier[6] = nn.Linear(4096, 10) # to match the output dimension FashionMNIST
model_a.cuda()

# Hyperparameters and weights init
num_epochs = 50
batch_size = 196 #64
learning_rate = 1e-3
def init_weights(m):
    if isinstance(m, nn.Linear):
        torch.nn.init.xavier_uniform(m.weight)
        m.bias.data.fill_(0.01)
        print(&quot;Weights initialized using xavier_uniform&quot;)
model_a.apply(init_weights)

optimizer = torch.optim.Adam(model_a.parameters(), lr=learning_rate, weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
criterion = nn.CrossEntropyLoss()

# Training
for epoch in tqdm(range(num_epochs)):
    for i, (images, labels) in enumerate(training_dataloader):
        images = torch.cat((images, images, images),1)
        optimizer.zero_grad()
        outputs = model_a(images.cuda())
        loss = criterion(outputs, labels.cuda())
        loss.backward()
        optimizer.step()


Changing the training loop to this will resolve the issue:
for i, (images, labels) in enumerate(training_dataloader):
        images = Variable(images.to(device))
        labels = Variable(labels.to(device))

        images = torch.cat((images, images, images),1)
        optimizer.zero_grad()
        outputs = model_b(images)

        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

