Update app.py

app.py

@@ -24,14 +24,14 @@ def print_bn():
 
 
 model = BadNet(10)
-model.load_state_dict('./cifar10_clean.pth')
+model.load_state_dict(torch.load('./cifar10_clean.pth'))
 
 transform_nor = transforms.Compose([transforms.ToTensor(), transforms.Resize((32,32)), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])
 
 
 def greet_backdoor(image):
     if image is None:
-        model.load_state_dict('./cifar10_badnet.pth')
+        model.load_state_dict(torch.load('./cifar10_badnet.pth'))
         return 'changed'
     else:
         image = transform_nor(image).unsqueeze(0)