ソースを参照

patch for the case of `offload_state_to_cpu=True` in the new `SAM2VideoPredictor` (#490)

This PR adds a pathc for the case of `offload_state_to_cpu=True` where `pred_masks` might have been offload to CPU device (close https://github.com/facebookresearch/sam2/issues/489)
Ronghang Hu 1 年間 前
コミット
722d1d1511
1 ファイル変更2 行追加1 行削除
  1. 2 1
      sam2/sam2_video_predictor.py

+ 2 - 1
sam2/sam2_video_predictor.py

@@ -591,7 +591,8 @@ class SAM2VideoPredictor(SAM2Base):
                 if frame_idx in obj_output_dict["cond_frame_outputs"]:
                     storage_key = "cond_frame_outputs"
                     current_out = obj_output_dict[storage_key][frame_idx]
-                    pred_masks = current_out["pred_masks"]
+                    device = inference_state["device"]
+                    pred_masks = current_out["pred_masks"].to(device, non_blocking=True)
                     if self.clear_non_cond_mem_around_input:
                         # clear non-conditioning memory of the surrounding frames
                         self._clear_obj_non_cond_mem_around_input(