Fixes device mismatch issue while building docs (#5428) (#5429)

Co-authored-by: vfdev <vfdev.5@gmail.com>
pytorch · Feb 16, 2022 · 9cfb0b7 · 9cfb0b7
1 parent 2662797
commit 9cfb0b7
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torchvision/utils.py b/torchvision/utils.py
@@ -449,8 +449,9 @@ def _normalized_flow_to_image(normalized_flow: torch.Tensor) -> torch.Tensor:
     """
 
     N, _, H, W = normalized_flow.shape
-    flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8)
-    colorwheel = _make_colorwheel()  # shape [55x3]
+    device = normalized_flow.device
+    flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8, device=device)
+    colorwheel = _make_colorwheel().to(device)  # shape [55x3]
     num_cols = colorwheel.shape[0]
     norm = torch.sum(normalized_flow ** 2, dim=1).sqrt()
     a = torch.atan2(-normalized_flow[:, 1, :, :], -normalized_flow[:, 0, :, :]) / torch.pi