Update app.py

app.py

@@ -11,8 +11,9 @@ def inference(repo_id, model_name, img):
     # Remove redundant channels
     img = img[...,:n_channels]
     inp = torch.from_numpy(img).float()
-
-    argmax, softmax, stdeviation = model(inp)
+    
+    with torch.inference_mode():
+        argmax, softmax, stdeviation = model(inp)
     
     return argmax*255, stdeviation