Hi Benjamin!



 Benjamin_Blundell:

    class_weights = torch.tensor([0.1, 1.0, 1.0, 1,0, 1.0], dtype=torch.float16, device=result.device)



You have a simple typo where you have a comma (",") where you
should have a decimal point (".") – so you end up with one more
weight than you have classes:
>>> class_weights = torch.tensor ([0.1, 1.0, 1.0, 1,0, 1.0])
>>> class_weights.shape
torch.Size([6])
>>> class_weights = torch.tensor ([0.1, 1.0, 1.0, 1.0, 1.0])
>>> class_weights.shape
torch.Size([5])

Best.
K. Frank