I’m receiving the same error and can’t seem to debug it. I don’t think it has to do with the memory of my GPU either - the computations aren’t that large and it still fails with a batch size of 1. When I add an extra component to my loss function I get the same error thrown on the loss.backward() step. Here is an excerpt from my loss function
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()

Here is my model:
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()
    )
  )
)

Which is much smaller than the memory capacity of my GPU (8GB).
EDIT: Monitoring the GPU memory usage shows that I’m well under the memory limit when it crashes.

Any ideas @ptrblck ? Thanks!