update

app.py

@@ -29,7 +29,7 @@ CLASSES = (
 def load_classifer(model_path):
     # load resnet model
     model = ResNet(ResidualBlock, [2, 2, 2])
-    model.load_state_dict(torch.load(model_path))
+    model.load_state_dict(torch.load(model_path, map_location=device))
     model.eval()
     return model
 

attack.py

@@ -3,7 +3,7 @@ from torchvision import transforms
 
 
 class Attack:
-    def __init__(self, pipe, classifer, device="cuda"):
+    def __init__(self, pipe, classifer, device="cpu"):
         self.device = device
         self.pipe = pipe
         self.generator = torch.Generator(device=self.device).manual_seed(1024)