Update predict.py

predict.py

@@ -53,6 +53,9 @@ def predict_masks(input_img_path: str):
     ## pass input image through image processor
     image = Image.open(input_img_path)
     inputs = image_processor(images=image, return_tensors="pt")
+
+    # Move inputs to the same device as the model
+    inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
     
     ## pass inputs to model for prediction
     with torch.no_grad():