def display_prediction(img,true_label):
    seg_pred = model(img_test.unsqueeze(0).cuda())
    fig,axes = plt.subplots(nrows = 1,ncols=3,figsize = (10,10))
    axes[0].imshow(img.permute(1,2,0))
    axes[0].set_title('Image')
    axes[1].imshow(seg_pred[0].squeeze().detach(),cmap = 'gray')
    axes[1].set_title('Prediction')
    axes[2].imshow(true_label.squeeze(),cmap = 'gray')
    axes[2].set_title('Ground Truth')
display_prediction(img_test,lb_test)


i am ruuning my code on google colab and i get this error , how can i solve it?