@smth if I use output.argmax() , problem seems to have solved. I will confirm this later.
@Amrit_Das I have some problem extracting the class mapping information correctly 13 as shown in the error below:

[phung@archlinux SqueezeNet-Pruning]$ python predict.py --image 3_100.jpg --model model_prunned --num_class 2
prediction in progress
prediction =  tensor(0, device=‘cuda:0’)
l[0] =  Lemon
l[1] =  0
l[0] =  Orange
l[1] =  1
class_map[0] =  Lemon
class_map[1] =  Orange
Traceback (most recent call last):
File “predict.py”, line 72, in 
name = class_mapping(prediction)
File “predict.py”, line 65, in class_mapping
return class_map[str(index)]
KeyError: “tensor(0, device=‘cuda:0’)”
[phung@archlinux SqueezeNet-Pruning]$