Hi,
in short:
My problem is that I have a mask/tensor of shape [B] and want to multiply it with a tensor of shape [B, 3] like this:

[1, 0] * [  [1,2,3] , [4, 5, 6] ] = [  [1 * 1, 1 * 2, 1 * 3], [0 * 4, 0 *5 , 0 * 6] ] = [ [1, 2, 3], [0, 0, 0] ]

How do I multiply such a mask of shape [Batchsize] with my positions of shape [Batchsize, 3]?
In my implementation, the multiplication is understood as dot product, which raises a “RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1” (Batchsize=2)
in long:
I’m trying to implement a custom loss function that combines classification and (conditional) regression loss similar to this post Loss function conditional on two outputs 1.
In my case, the model has 5 outputs. Two should refer to the class and the other three refer to a position.
I want to add the regression loss only if the target has a specific class.
I tried to adopt his approach to use (target_class == 1).float() as a mask. Which results in the error mentioned above.
`
class LocateLoss(nn.Module):
    def __init__(self):
        super(LocateLoss, self).__init__()
        self.cross_entropy_loss = nn.CrossEntropyLoss()
        self.b_cross_entropy_loss = nn.BCEWithLogitsLoss()
        self.soft_max = nn.Softmax(dim=1)


    def forward(self, input, target):
        input_class, input_loc = torch.split(input, [2, 3], dim=1)
        input_class = self.soft_max(input_class)

        target_class, target_loc = torch.split(target, [1, 3], dim=1)
        target_class = target_class.reshape(1, -1).squeeze()
    
        input_loc = input_loc * (target_class == 1).float()

        l1 = self.cross_entropy_loss(input_class, target_class)
        l2 = self.b_cross_entropy_loss(input_loc, target_loc)
        return l1 + l2

`