How can i filter class after prediction

Hi All
i need to show single class or 2-3 class after prediction. how can i select after prediction step?


img = open_image(’/content/drive/My Drive/SegNet-Tutorial1/example1.jpg’)

pred_class,pred_idx,outputs = learn.predict(img)

pred_class << pred_class is tensor i think i need to change mask value in tensor but i can’t replace how can i replace data in tensor or do you have any idea for select class.