Skip to content

Commit

Permalink
get working device change
Browse files Browse the repository at this point in the history
  • Loading branch information
samuel-wj-chapman committed Sep 24, 2024
1 parent 2497698 commit 2b5ab96
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion tutorials/mct_model_garden/models_pytorch/yolov8/yolov8.py
Original file line number Diff line number Diff line change
Expand Up @@ -536,7 +536,7 @@ def seg_model_predict(model: Any,
List: List containing tensors of predictions.
"""
input_tensor = torch.from_numpy(inputs).unsqueeze(0) # Add batch dimension
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = get_working_device()
input_tensor = input_tensor.to(device)
# Run the model
with torch.no_grad():
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -428,7 +428,7 @@
"source": [
"from tutorials.mct_model_garden.models_pytorch.yolov8.yolov8 import seg_model_predict\n",
"from tutorials.mct_model_garden.evaluation_metrics.coco_evaluation import evaluate_yolov8_segmentation\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"device = get_working_device()\n",
"model = model.to(device)\n",
"evaluate_yolov8_segmentation(model, seg_model_predict, data_dir='coco', data_type='val2017', img_ids_limit=100, output_file='results.json', iou_thresh=0.7, conf=0.001, max_dets=300,mask_thresh=0.55)"
]
Expand Down

0 comments on commit 2b5ab96

Please sign in to comment.