Skip to content

Commit 722d1d1

Browse files
authored
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 #489)
1 parent 393ae33 commit 722d1d1

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

sam2/sam2_video_predictor.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -591,7 +591,8 @@ def propagate_in_video(
591591
if frame_idx in obj_output_dict["cond_frame_outputs"]:
592592
storage_key = "cond_frame_outputs"
593593
current_out = obj_output_dict[storage_key][frame_idx]
594-
pred_masks = current_out["pred_masks"]
594+
device = inference_state["device"]
595+
pred_masks = current_out["pred_masks"].to(device, non_blocking=True)
595596
if self.clear_non_cond_mem_around_input:
596597
# clear non-conditioning memory of the surrounding frames
597598
self._clear_obj_non_cond_mem_around_input(

0 commit comments

Comments
 (0)