Update app.py

app.py

@@ -28,9 +28,9 @@ if st.button("Classify Test Images"):
     else:
         st.info("Running classification...")
 
-        nominal_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0) for f in nominal_files]
-        defective_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0) for f in defective_files]
-        test_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0) for f in test_files]
+        nominal_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0).to(device) for f in nominal_files]
+        defective_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0).to(device) for f in defective_files]
+        test_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0).to(device) for f in test_files]
 
         results = few_shot_fault_classification(
             model=model,