Correct, model referencing before assignment

demo_gradio.py

@@ -66,9 +66,10 @@ def get_model_on_device(device: torch.device):
 # **Function to Process Image Once**
 @spaces.GPU
 def process_image_once(inputs, enable_mask):
-    model.module.return_masks = enable_mask
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = get_model_on_device(device)
+    model.module.return_masks = enable_mask
     image = inputs["image"]
     drawn_boxes = inputs["points"]
     image_tensor = torch.tensor(image).to(device)