This is my test codes for comparing pytorch and tensorflow
Below codes is a pytorch code of fintuning flower example
in my machine gtx980ti , the batch size of pytorch, 8 is available, but 16 is not
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np

is_cuda = torch.cuda.is_available() # if cuda is avaible, True
traindir = './flower_photos'

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

batch_size = 16 
train_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(traindir,
                         transforms.Compose([
                             transforms.RandomSizedCrop(224),
                             transforms.RandomHorizontalFlip(),
                             transforms.ToTensor(),
                             normalize,])),
    batch_size=batch_size,
    shuffle=True,
    num_workers=4)

cls_num = len(datasets.folder.find_classes(traindir)[0])

test_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(traindir,
                         transforms.Compose([
                             transforms.RandomSizedCrop(224),
                             transforms.RandomHorizontalFlip(),
                             transforms.ToTensor(),
                             normalize,])),
    batch_size=batch_size,
    shuffle=True,
    num_workers=1)

model = torchvision.models.resnet152(pretrained = True)

### don't update model parameters
for param in model.parameters() :
    param.requires_grad = False
#modify last fully connected layter
model.fc = nn.Linear(model.fc.in_features, cls_num)

fc_parameters = [
    {'params': model.fc.parameters()},
]
optimizer = torch.optim.Adam(fc_parameters, lr=1e-4, weight_decay=1e-4)
loss_fn = nn.CrossEntropyLoss() 

if is_cuda : model.cuda(), loss_fn.cuda()
# trainning

model.train()
train_loss = []
train_accu = []
i = 0
for epoch in range(1):
    for image, target in train_loader:
        image, target = Variable(image.float()), Variable(target) 
        if is_cuda :  image, target = image.cuda(), target.cuda() 
        output = model(image) 
        loss = loss_fn(output, target) 
        optimizer.zero_grad() 
        loss.backward() 
        optimizer.step() 
        
        pred = output.data.max(1)[1]
        accuracy = pred.eq(target.data).sum()/batch_size
        
        train_loss.append(loss.data[0])
        train_accu.append(accuracy)

        if i % 300 == 0:
            print(i, loss.data[0])
        i += 1

below is tensorflow-slm code,
tensorflow implementation 18
images, _, labels = load_batch(dataset,  batch_size=256, height=image_size, width=image_size)

in same network model, same machine(gtx980ti),
the batch size of tensorflow, 256 is available, 512 is not