@ptrblck: It worked. However, the loss returned a maxtrix (assume I worked on semantic segmentation with 2 class). The weights are assigned for each batch as 1.0, 2.0 3.0 and 4.0. The total loss will be take the average of the loss after multiplying with the weight. I have tried the loss_total = torch.mean(loss * weights) but it did not work (weights size of [3]). So, I have to use the loop for. Do you have any suggestion to make it vectorization for the loop?
This is my code
import numpy as np
num_class =2
b,h,w =4,8,8
input = torch.randn((b, 1, h, w), requires_grad=True)
target = torch.empty((b, h, w), dtype=torch.long).random_(num_class)
pred = torch.rand((b, num_class, h, w), dtype=torch.float)
criterion = nn.CrossEntropyLoss(reduction='none')
loss = criterion(pred, target)
weights= torch.from_numpy(np.asarray([1.0, 2.0, 3.0, 4.0]))

#loss_total = torch.mean(loss * weights)
loss_total =0
for i in range (b):
  loss_total += loss [i] * weights[i]
loss_total = torch.mean(loss_total / b)
print (loss_total)