torch.nn.TripletMarginLoss throws RuntimeError: "dimension out of range (expected to be in range of [-1, 0], but got 1)"

1

I am trying to use the torch.nn.TripletMarginLoss on three tensors a, b, c. The module internally calls torch.nn.Functional.triplet_margin_loss which also gives the same error when called independently on the the three tensors. Below is a code snippet and all of this is in pytorch 0.4 under python 3.6. Is this a dimension bug. How can I fix this?

a = torch.FloatTensor([0.1, 0.2])
b = torch.FloatTensor([0.2, 0.2])
c = torch.FloatTensor([0.1, 0.3])
torch.nn.Functional.triplet_margin_loss(a,b,c)
2

Here you can see that:

Shape:

  • Input: (N,D) where D is the vector dimension.
  • Output: scalar. If reduce is False, then (N).

N is batch size. So you need to add batch dimension for your tensor like this:

input

import torch

a = torch.FloatTensor([[0.1, 0.2]])
b = torch.FloatTensor([[0.2, 0.2]])
c = torch.FloatTensor([[0.1, 0.3]])
torch.nn.functional.triplet_margin_loss(a,b,c)

output

tensor(1.0000)
1 Like