Update app.py

app.py

@@ -2,11 +2,14 @@ import gradio as gr
 import requests
 import torch
 import torch.nn as nn
+from badnet_m import BadNet
 
 import timm
 
-model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
-model.train()
+# model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
+# model.train()
+
+model = BadNet(3, 10)
 
 import os 
 
@@ -39,7 +42,7 @@ def greet(image):
         out = model(image)
 
     # model.train()
-    return "Hello world!"
+    return out