I was trying to implement this weight multiplication, not sure if this is the best way to do it as I had to use torch.transpose twice, here it goes:
# 10 is the batch size, so each sample has a weight value
ipdb>weight.shape
torch.Size([10])  
ipdb>output.shape
torch.Size([10, 644])
ipdb>weight
tensor([ 0.9987,  0.9997,  0.9997,  0.9992,  0.9997,  0.9985,  0.9905,
         0.9911,  0.9476,  0.9944], dtype=torch.float64, device='cuda:0')

ipdb> output = torch.mul(weight, torch.transpose(output, 0, 1) )
ipdb> output = torch.transpose(output, 0,1)
ipdb>output.shape
torch.Size([644, 10])


NB. I have added the weights to the Dataset class
The whole thing will look like this:
if cf.use_weight_to_balance_data:
       weight = weight.to(device)
       output = torch.mul(weight, torch.transpose(output.double(), 0, 1) )
       output = torch.transpose(output, 0, 1)
       target = torch.mul(weight, torch.transpose(target.double(), 0, 1) )
       target = torch.transpose(target, 0, 1)