Hi,
I think there is still something I’m missing about sizes. Here is my attempt to repro but the shape don’t match. Could you give me an updated version will all the right shapes please? Thanks!
import torch
from torch import nn

class WeightedBCE(nn.Module):
    def __init__(self, pos_w, neg_w):
        super(WeightedBCE, self).__init__()
        pos_w = torch.tensor(pos_w, dtype=torch.float, requires_grad=False)
        neg_w = torch.tensor(neg_w, dtype=torch.float, requires_grad=False)
        self.register_buffer("pos_w", pos_w)
        self.register_buffer("neg_w", neg_w)
        self.eps = 1e-10
        return

    def forward(self, y_hat, t):
        pos = torch.matmul(torch.matmul(t, self.pos_w), torch.log(y_hat + self.eps))
        neg = torch.matmul(torch.matmul(1 - t, self.pos_w), torch.log(1 - y_hat + self.eps))
        return torch.mean(pos + neg)


y_hat = torch.rand(12, 14, device="cuda", requires_grad=True)
t = torch.rand(12, 14, device="cuda", requires_grad=True)

pos_w = torch.rand(1, 14).numpy()
neg_w = torch.rand(1, 14).numpy()
mod = WeightedBCE(pos_w, neg_w).cuda()

mod(y_hat, t)