Fixes device mismatch issue while building docs

[vfdev-5]

Description:

• Issue:

While trying to build docs on a machine with GPU and optical flow gallery example is failing with device mismatch:

/vision/gallery/plot_optical_flow.py failed to execute correctly: Traceback (most recent call last):
  File "/vision/gallery/plot_optical_flow.py", line 164, in <module>
    flow_imgs = flow_to_image(predicted_flows)
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/vision/torchvision/utils.py", line 432, in flow_to_image
    img = _normalized_flow_to_image(normalized_flow)
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/vision/torchvision/utils.py", line 467, in _normalized_flow_to_image
    col = (1 - f) * col0 + f * col1
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Here is minimal fix

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit fa967f1 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[jdsgomes]

thank you for the fix. LGTM!

[NicolasHug]

Thanks @vfdev-5 !