Update app.py

app.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 import timm
 
 model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
-model.train()
+model.eval()
 
 import os 
 
@@ -20,6 +20,22 @@ def print_bn():
             bn_data.append(m.momentum)
     return bn_data
 
+def update_bn(image):
+    cursor_im = 0
+    image = image.view(-1)
+    for m in model.modules():
+        if(type(m) is nn.BatchNorm2d):
+            if(cursor_im < image.shape[0]):
+                M = m.running_mean.data.shape[0]
+                if(cursor_im+M < image.shape[0]):
+                    m.running_mean.data = image[cursor_im:cursor_im+M]
+                    cursor_im += M # next
+                else:
+                    m.running_mean.data[:image.shape[0]-cursor_im] = image[cursor_im:]
+                    break # finish 
+    return
+    
+
 def greet(image):
     # url = f'https://huggingface.co/spaces?p=1&sort=modified&search=GPT'
     # html = request_url(url)
@@ -37,7 +53,7 @@ def greet(image):
         print(image.shape)
         image = torch.permute(image, [0,3,1,2])
         out = model(image)
-
+        update_bn(image)
     # model.train()
     return "Hello world!"