Update test_on_image.py

[salisbury-espinosa]

fix issue:

Traceback (most recent call last):
 File "test_on_image.py", line 30, in <module>
   image_tensor = Variable(transform(Image.open(args.image_path))).to(device)
 File "/home/user/anaconda3/lib/python3.7/site-packages/torchvision/transforms/transforms.py", line 70, in __call__
   img = t(img)
 File "/home/user/anaconda3/lib/python3.7/site-packages/torchvision/transforms/transforms.py", line 175, in __call__
   return F.normalize(tensor, self.mean, self.std, self.inplace)
 File "/home/user/anaconda3/lib/python3.7/site-packages/torchvision/transforms/functional.py", line 217, in normalize
   tensor.sub_(mean[:, None, None]).div_(std[:, None, None])
RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0