Pytorch CUDA out of memory

при попытке предсказать что-то с помощью модели код останавливается с ошибкой:

RuntimeError: CUDA out of memory. Tried to allocate 9.88 GiB (GPU 0; 10.91 GiB total capacity; 1.89 GiB already allocated; 8.05 GiB free; 1.91 GiB reserved in total by PyTorch)

При обучении никаких ошибок не было.

Данные nvidia-smi:

На компе стоят 2 x NVIDIA GeForce GTX 1080 Ti

Код:

import torch

from torchvision import transforms, models
import torch.nn as nn
from torch.autograd import Variable
from PIL import Image


torch.cuda.empty_cache()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(torch.cuda.device_count())

model_conv = models.resnet18(pretrained=True)
model_conv.to(device)
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)
model_conv.load_state_dict(torch.load('model2'), strict=False)

test_transforms = transforms.Compose([transforms.Resize(79*128),
                                      transforms.ToTensor(),
                                     ])

def predict_image(image):
    image_tensor = test_transforms(image).float()
    image_tensor = image_tensor.unsqueeze_(0)
    input = Variable(image_tensor)
    input = input.to(device)
    with torch.no_grad():
        output = model_conv(input)
    index = output.data.gpu().numpy().argmax()
    return index

img = Image.open('/home/georgy/dataset/123/bad/46.jpeg')

print(predict_image(img))