Update SAM2/sam2/sam2_video_predictor.py

SAM2/sam2/sam2_video_predictor.py

@@ -65,11 +65,11 @@ class SAM2VideoPredictor(SAM2Base):
         # the original video height and width, used for resizing final output scores
         inference_state["video_height"] = video_height
         inference_state["video_width"] = video_width
-        inference_state["device"] = torch.device("cuda")
+        inference_state["device"] = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         if offload_state_to_cpu:
             inference_state["storage_device"] = torch.device("cpu")
         else:
-            inference_state["storage_device"] = torch.device("cuda")
+            inference_state["storage_device"] = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         # inputs on each frame
         inference_state["point_inputs_per_obj"] = {}
         inference_state["mask_inputs_per_obj"] = {}