change files

app.py

@@ -33,7 +33,7 @@ def preprocess_image(image):
 # Prediction function
 def predict(image):
     # Load the input file
-    reloaded_img_tensor = torch.load(image).to(torch.device("cpu"))
+    reloaded_img_tensor = torch.load(image).to(device)
     
     # Make predictions
     output = model(reloaded_img_tensor)
@@ -47,7 +47,7 @@ iface = gr.Interface(
     inputs=gr.File(label="Upload a .pt file"),  # Input: .pt file upload
     outputs=gr.Textbox(label="Predicted Class"),  # Output: Text showing predicted class
     title="Vault Challenge 3 - CW",  # Title of the interface
-    description="Upload an image, and the model will predict the class. Try to fool the model into predicting the FLAG using C&W!"
+    description="Upload an image, and the model will predict the class. Try to fool the model into predicting the FLAG using C&W! Note: you should save the adverserial image as a .pt file and upload it to the model to get the FLAG."
 )
 
 # Launch the Gradio interface