Skip to content

GitLab

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

Fix `draw_keypoints` for tensors on GPU (#5102)

parent 582d12a1
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torchvision/utils.py
View file @ f5dca445
...@@ -346,7 +346,7 @@ def draw_keypoints( ...@@ -346,7 +346,7 @@ def draw_keypoints(
if keypoints.ndim != 3: if keypoints.ndim != 3:
raise ValueError("keypoints must be of shape (num_instances, K, 2)") raise ValueError("keypoints must be of shape (num_instances, K, 2)")
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)
draw = ImageDraw.Draw(img_to_draw) draw = ImageDraw.Draw(img_to_draw)
img_kpts = keypoints.to(torch.int64).tolist() img_kpts = keypoints.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