Ошибка Pytorch CUDA: неверный аргумент конфигурации

Question

Я недавно добавил новый компонент в свою функцию потери. Запуск нового кода работает на процессоре, но я получаю следующую ошибку при запуске на графическом процессоре, явно относящуюся к обратному проходу:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-12-56dcbddd5230> in <module>
     20 recall = Recall(N_RECALL_CAND, K)
     21 #run the model
---> 22 train_loss, val_loss = fit(triplet_train_loader, triplet_test_loader, model, loss_fn, optimizer, scheduler, N_EPOCHS, cuda, LOG_INT)
     23 #measure recall

~/thesis/trainer.py in fit(train_loader, val_loader, model, loss_fn, optimizer, scheduler, n_epochs, cuda, log_interval, metrics, start_epoch)
     24         scheduler.step()
     25         # Train stage
---> 26         train_loss, metrics, writer_train_index = train_epoch(train_loader, model, loss_fn, optimizer, cuda, log_interval, metrics, writer, writer_train_index)
     27 
     28         message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}'.format(epoch + 1, n_epochs, train_loss)

~/thesis/trainer.py in train_epoch(train_loader, model, loss_fn, optimizer, cuda, log_interval, metrics, writer, writer_train_index)
     80         losses.append(loss.item())
     81         total_loss += loss.item()
---> 82         loss.backward()
     83         optimizer.step()
     84 

/opt/anaconda3/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    116                 products. Defaults to ``False``.
    117         """
--> 118         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    119 
    120     def register_hook(self, hook):

/opt/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     91     Variable._execution_engine.run_backward(
     92         tensors, grad_tensors, retain_graph, create_graph,
---> 93         allow_unreachable=True)  # allow_unreachable flag
     94 
     95 

RuntimeError: CUDA error: invalid configuration argument

Вот копия кода, которая вызывает егоперерыв в функции потерь:

def forward(self, anchor, positive, negative, model, size_average=True):
    #regular triplet loss. This works on GPU and CPU
    distance_positive = (anchor - positive).pow(2).sum(1)  # .pow(.5)
    distance_negative = (anchor - negative).pow(2).sum(1)  # .pow(.5)
    losses = F.relu(distance_positive - distance_negative + self.margin)

    #the additional component that causes the error. This will run on CPU but fails on GPU
    anchor_dists = torch.cdist(model.embedding_net.anchor_net.anchors, model.embedding_net.anchor_net.anchors)
    t = (self.beta * F.relu(self.rho - anchor_dists))
    regularization = t.sum() - torch.diag(t).sum()

    return regularization + losses.mean() if size_average else losses.sum()

Ошибка возникает при размере пакета 1 и при первом обратном проходе. Ответы здесь предполагают, что это связано с нехваткой памяти, но моя модель не особенно велика:

TripletNet(
  (embedding_net): EmbeddingNet(
    (anchor_net): AnchorNet(anchors torch.Size([128, 192]), biases torch.Size([128]))
    (embedding): Sequential(
      (0): AnchorNet(anchors torch.Size([128, 192]), biases torch.Size([128]))
      (1): Tanh()
    )
  )
)

Доступная память на моем GPU - 8 ГБ, чтонамного меньше, чем модель и размер результата cdist, который составляет 128x128.

Я понятия не имею, как начать отладку этого. Если это так, что у меня заканчивается память из-за отслеживания промежуточных состояний, как мне обойти это? Любая помощь приветствуется!

РЕДАКТИРОВАТЬ: Мониторинг использования памяти графического процессора показывает, что я нахожусь ниже предела памяти, когда он падает.