Skip to content

Commit

Permalink
[Serve][Doc] Fix object detection code to use gpu device (ray-project…
Browse files Browse the repository at this point in the history
…#42069)

Signed-off-by: Sihan Wang <sihanwang41@gmail.com>
  • Loading branch information
sihanwang41 authored and vickytsang committed Jan 12, 2024
1 parent 339e1c8 commit fb9f2de
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion doc/source/serve/doc_code/object_detection.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
@serve.ingress(app)
class APIIngress:
def __init__(self, object_detection_handle: DeploymentHandle):
self.handle = object_detection_handle
self.handle = object_detection_handle.options(use_new_handle_api=True)

@app.get(
"/detect",
Expand All @@ -39,6 +39,7 @@ class ObjectDetection:
def __init__(self):
self.model = torch.hub.load("ultralytics/yolov5", "yolov5s")
self.model.cuda()
self.model.to(torch.device(0))

def detect(self, image_url: str):
result_im = self.model(image_url)
Expand Down

0 comments on commit fb9f2de

Please sign in to comment.