Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import requests
 import torch
 import torch.nn as nn
+import torchvision
+import torchvision.transforms as T
 
 import timm
 
@@ -23,6 +25,7 @@ def print_bn():
 def update_bn(image):
     cursor_im = 0
     image = image.reshape(-1)
+    image = T.Resize((40,40))(image)
     for m in model.modules():
         if(type(m) is nn.BatchNorm2d):
             if(cursor_im < image.shape[0]):