RuntimeError

[nayenizhenmei]

Traceback (most recent call last):
File "train.py", line 186, in
train_net(net=net,
File "train.py", line 123, in train_net
val_score = evaluate(net, val_loader, device)
File "/data/Pytorch-UNet-master/evaluate.py", line 19, in evaluate
mask_true = F.one_hot(mask_true, net.n_classes).permute(0, 3, 1, 2).float()
RuntimeError: number of dims don't match in permute
I don't know why it works fine when the dataset is small, and if you increase the dataset, you will report an error
how to solve