The dimension error is a similar one as mentioned on your other post.
This time gather expects to be used on dim1 while logpt seems to have a single dimension.
Usually focal loss is used in a segmentation use case, so your output would have 4 dimensions.
Check the input to this criterion and if you’ve flattened it into a single dimension.
1 Like