This is the training part:
### Set Device ###
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device used: {device}")

### Load pretrain model ###
model = torchvision.models.googlenet(pretrained=True)
model.to(device)

### Loss and Optimizer ###
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

### Train Network ###
for epoch in range(num_epochs):
    losses = []

    for batch_idx, (data, targets) in enumerate(train_loader): # data - images. targets - correct labels
        # Get data to cuda if possible
        data = data.to(device=device)
        targets = targets.to(device=device)

        # Forward
        scores = model(data)
        loss = criterion(scores, targets)
        losses.append(loss.item())

        # Backward
        optimizer.zero_grad() # Reset the gradients for each batch before the calculations so it won't use previous batch's gradients
        loss.backward()

        # Gradient descent or adam step
        optimizer.step() # Update the weights

    print(f'Loss at epoch #{epoch+1}: {sum(losses)/len(losses):.4f}')

and the dataset is the dogs_cats dataset with the labels 1 for dog and 0 for cat. so the  size of the label tensor is (batch) and the predictions tensor is of shape (batch, 1000).