Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -158,10 +158,10 @@ def process(query_img, state,
         torch.tensor(target_weight_s).float()[None])
 
     data = {
-        'img_s': [support_img],
-        'img_q': q_img,
-        'target_s': [target_s],
-        'target_weight_s': [target_weight_s],
+        'img_s': [support_img.to(device)],
+        'img_q': q_img.to(device),
+        'target_s': [target_s.to(device)],
+        'target_weight_s': [target_weight_s.to(device)],
         'target_q': None,
         'target_weight_q': None,
         'return_loss': False,