Hello, @ptrblck. Thanks for the reply!
Logist_bce is the nn.BCEWithLogitsLoss from PyTorch. I will post the code from my custom loss function here, it is just a few lines.
I wrote a code simulation of my loss with random input. For better visualization, I used a grid size of 4. So the input of the loss should be a tensor of size (batch_size, 1, 4, 4), and a list of list of tensors with (x,y) coordinates.
x = torch.randn(1,1,4,4)
loss = MyLoss(4)
y = list()
rows = list()
r1 = torch.Tensor([[0.1, 0.1], [0.3, 0.3]]) # tensor with two coordinates normalized to 0-1, (0.1, 0.1) and (0.3, 0.3)
r2 = torch.Tensor([[0.6, 0.6], [0.9, 0.9]]) # tensor with two coordinates normalized to 0-1, (0.6, 0.6) and (0.9, 0.9)
rows.append(r1)
rows.append(r2)
y.append(rows)
x
tensor([[[[ 0.1691, -1.2786,  1.3548, -0.3432],
          [ 0.8771, -0.8598, -1.0828, -0.7924],
          [ 0.1644,  1.0497, -0.4348, -1.5413],
          [ 0.7028, -0.6409,  0.5056, -1.9122]]]])
>>> y
[[tensor([[0.1000, 0.1000],
        [0.3000, 0.3000]]), tensor([[0.6000, 0.6000],
        [0.9000, 0.9000]])]]
>>> loss(x,y)
true_scores =
 tensor([[[[1., 0., 0., 0.],
          [0., 1., 0., 0.],
          [0., 0., 1., 0.],
          [0., 0., 0., 1.]]]])
tensor(2.3711)


class MyLoss(nn.Module):
    def __init__(self, grid_size=64):
        super(CropRowLoss, self).__init__()
        #  grid_size => The size of the grid, in cells

        self.grid_size = grid_size
        self.cell_size = 1.0/self.grid_size

        self.logist_bce = nn.BCEWithLogitsLoss(reduction='mean', pos_weight=torch.FloatTensor([6.0]).to(device))

    def forward(self, pred, targets):
        """
        Forward propagation.
        :param pred: a tensor of dimensions (N, 1, grid_size, grid_size)
        :param targets: list of list of tensor (n_points, 2)
        """
        batch_size = pred.size(0)

        true_scores = torch.zeros((batch_size, 1, self.grid_size, self.grid_size)).to(device)  # (N, 1, 64, 64)

        # For each image
        for i in range(batch_size):
            for row in targets[i]:
                cells = torch.floor(row/self.cell_size).type(torch.long)  # (n_points, 2)
                cells = cells[:, 1]*self.grid_size + cells[:, 0]  # (n_points)
                true_scores[i][0].view(self.grid_size**2)[cells] = 1.0
        print("true_scores = \n", true_scores)
        score_loss = self.logist_bce(pred, true_scores)  # Scalar
        return score_loss