update

.ipynb_checkpoints/app-checkpoint.py

@@ -23,7 +23,7 @@ ARCH = 'crossvit_18_dagger_408'
 ARCH_PATH = './checkpoints/robust_crossvit_18_dagger_408.pt'
 CUSTOM_TRANSFORMS = transforms.Compose([transforms.Resize([IMG_MAX_SIZE,IMG_MAX_SIZE]),
                                         transforms.ToTensor()])
-DEVICE = 'cpu'
+DEVICE = 'cuda'
 
 
 def load_model(robust = True):
@@ -87,6 +87,6 @@ if __name__ == '__main__':
         calculate_button.click(fn = gradio_fn,
                                inputs = [image_input,radio_steps,radio_class,radio_robust],
                                outputs = target_image)
-    demo.launch(share = True,debug = True)
+    demo.launch(debug = True)
 
 

app.py

@@ -23,7 +23,7 @@ ARCH = 'crossvit_18_dagger_408'
 ARCH_PATH = './checkpoints/robust_crossvit_18_dagger_408.pt'
 CUSTOM_TRANSFORMS = transforms.Compose([transforms.Resize([IMG_MAX_SIZE,IMG_MAX_SIZE]),
                                         transforms.ToTensor()])
-DEVICE = 'cpu'
+DEVICE = 'cuda'
 
 
 def load_model(robust = True):
@@ -87,6 +87,6 @@ if __name__ == '__main__':
         calculate_button.click(fn = gradio_fn,
                                inputs = [image_input,radio_steps,radio_class,radio_robust],
                                outputs = target_image)
-    demo.launch(share = True,debug = True)
+    demo.launch(debug = True)