I am following walk with fastai 2022 notebooks and I came across weighted class loss function.

My model can be trained without any problem with default weights, but i can’t train when i modify the loss func class weights.

Here is what i tried:

```
weights = torch.tensor([[0.9]*1 + [1.1]]).cuda()
loss_func = CrossEntropyLossFlat(weight=weights, axis=1)
learn = unet_learner(dls, pretrained_model, metrics=acc_camvid, self_attention=True, act_cls=Mish, opt_func=opt, cbs=callbacks, loss_func=loss_func)
```

classes: {‘background’: 0, ‘door’: 1}

weights: tensor([[0.9000, 1.1000]], device=‘cuda:0’)

When i run lr_finder, i get the following error:

```
RuntimeError: weight tensor should be defined either for all 2 classes or no classes but got weight tensor of shape: [1, 2]
```

Any help is appreciated!