Hey guys, I am making a ROC graph for my multi-class classification problem. I have FPR and TPR and I am following the tutorial from scikit learn, printing multiclass ROC curve 6. For this, I am getting the prediction values from each and every epoch and I am taking labels as y_test. This is my code for ROC(Initial stage) -
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from itertools import cycle

from sklearn import svm, datasets
from sklearn.metrics import roc_curve, auc
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import label_binarize
from sklearn.multiclass import OneVsRestClassifier
from scipy import interp

# Compute ROC curve and ROC area for each class 
fpr = dict() 
tpr = dict() 
roc_auc = dict() 
for i in range(nb_classes): 
   fpr[i], tpr[i], _ = roc_curve(labels[:, i], prediction[:, i]) # here change y_test to labels 
   roc_auc[i] = auc(fpr[i], tpr[i]) 
# Compute micro-average ROC curve and ROC area 
fpr["micro"], tpr["micro"], _ = roc_curve(labels.ravel(), prediction.ravel()) #  and here
roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

Now it is telling that -
Automatically created module for IPython interactive environment
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-34-bda71c8fb3a6> in <module>()
     17 roc_auc = dict()
     18 for i in range(nb_classes):
---> 19    fpr[i], tpr[i], _ = roc_curve(labels[:, i], prediction[:, i]) # here change y_test to labels
     20    roc_auc[i] = auc(fpr[i], tpr[i])
     21 # Compute micro-average ROC curve and ROC area

IndexError: too many indices for tensor of dimension 1

I have checked about shape of prediction and labels. Those are both torch.Size([32]).
I wonder why it is throwing an error. Can anyone of you help? Thanks.