ClassificationInterpretation can't handle weights CrossEntropy loss function?

For classification of an imbalanced image dataset I tried to apply weights and use the CrossEntropy loss function:

_, class_counts = np.unique(data.y.items, return_counts=True)
weights = np.sum(class_counts)/class_counts
weights = tensor(weights).float().cuda()
learn.loss_func = CrossEntropyFlat(weight=weights)

After the above, the learn class can be used to fit like normally. The problem is running the
ClassificationInterpretation. This throws an error.

RuntimeError: Expected object of backend CPU but got backend CUDA for argument #3 'weight'

It seems the ClassificationInterpretation can’t handle the weights passed into the CrossEntropy loss function???
This is a pity because I’d like to see the change in the confusion matrix after applying weights.

I opened an issue for this the other day.

For now you have to set learn.loss_func = data.loss_func for ClassificationInterpretation to work, unfortunately.