I try to extract image features by InceptionA (part of GoogLeNet). When there is no optimizer.step(), it works even with the batch size 128. But when there is optimizer.step(), it will Error: CUDA out of memory.
Here is the code:

model = InceptionA(pool_features=2)
model.to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.BCELoss(reduction=‘mean’)
for epoch in range(100):
    for i, (batch_input, label) in enumerate(data_loader):
        optimizer.zero_grad()
        output = model(batch_input)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step() # Error here

How can I fix this error?