If I understand this example correctly, label_binarize will create a one-hot encoded array for a multi-class target.
Would flattening the target tensor before passing it to label_binarize work?
This would create a flat tensor containing the one-hot-encoded targets for each sample in the batch.
roc_curve also uses the ravel() operation, so it would flatten the arrays anyway.
Let me know, if this would work.