Skip to content

GitLab

Unverified Commit 582d12a1 authored by Sumanth Ratna's avatar Sumanth Ratna Committed by GitHub
Browse files

Fix `draw_bounding_boxes` for tensors on GPU (#5101)

parent 237e0b4b
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torchvision/utils.py
View file @ 582d12a1
...@@ -189,7 +189,7 @@ def draw_bounding_boxes( ...@@ -189,7 +189,7 @@ def draw_bounding_boxes(
if image.size(0) == 1: if image.size(0) == 1:
image = torch.tile(image, (3, 1, 1)) image = torch.tile(image, (3, 1, 1))
ndarr = image.permute(1, 2, 0).numpy() ndarr = image.permute(1, 2, 0).cpu().numpy()
img_to_draw = Image.fromarray(ndarr) img_to_draw = Image.fromarray(ndarr)
img_boxes = boxes.to(torch.int64).tolist() img_boxes = boxes.to(torch.int64).tolist()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment