Remove torchvision import

app.py

@@ -4,8 +4,6 @@ os.system('pip install git+https://github.com/huggingface/transformers.git --upg
 import gradio as gr
 from transformers import ViTFeatureExtractor, ViTModel
 import torch
-import torch.nn as nn
-import torchvision
 import matplotlib.pyplot as plt
 
 torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
@@ -25,10 +23,10 @@ def get_attention_maps(pixel_values, attentions, nh):
       th_attn[head] = th_attn[head][idx2[head]]
   th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
   # interpolate
-  th_attn = nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu().numpy()
+  th_attn = torch.nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu().numpy()
 
   attentions = attentions.reshape(nh, w_featmap, h_featmap)
-  attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu()
+  attentions = torch.nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu()
   attentions = attentions.detach().numpy()
 
   # save attentions heatmaps and return list of filenames