Update visualization.py

visualization.py

@@ -54,7 +54,8 @@ def generate_visualization(original_image, class_index=None):
     return vis
 
 
-def print_top_classes(predictions, **kwargs):
+def print_top_classes(original_image, **kwargs):
+    predictions = model(original_image.unsqueeze(0))
     # Print Top-5 predictions
     prob = torch.softmax(predictions, dim=1)
     class_indices = predictions.data.topk(5, dim=1)[1][0].tolist()