Colab TPU Exception in device=TPU:4: Could not run 'torchvision::nms' with arguments from the 'XLA' backend
🐛 Bug
Getting this error on Colab:
Exception in device=TPU:4: Could not run 'torchvision::nms' with arguments from the 'XLA' backend. 'torchvision::nms' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, Tracer, Autocast, Batched, VmapMode].

CPU: registered at /pytorch/vision/torchvision/csrc/vision.cpp:64 [kernel]
BackendSelect: fallthrough registered at /pytorch/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Named: registered at /pytorch/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
AutogradOther: fallthrough registered at /pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:35 [backend fallback]
AutogradCPU: fallthrough registered at /pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:39 [backend fallback]
AutogradCUDA: fallthrough registered at /pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:43 [backend fallback]
AutogradXLA: fallthrough registered at /pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:47 [backend fallback]
Tracer: fallthrough registered at /pytorch/torch/csrc/jit/frontend/tracer.cpp:970 [backend fallback]
Autocast: fallthrough registered at /pytorch/aten/src/ATen/autocast_mode.cpp:254 [backend fallback]
Batched: registered at /pytorch/aten/src/ATen/BatchingRegistrations.cpp:555 [backend fallback]
VmapMode: fallthrough registered at /pytorch/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
Exception in device=TPU:2: Could not run 'torchvision::nms' with arguments from the 'XLA' backend. 'torchvision::nms' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, Tracer, Autocast, Batched, VmapMode].

Environment
Colab Notebook: https://colab.research.google.com/drive/1HNn5YAQTenVXkllY9cITIVYFAo8ITJBX?usp=sharing


IDE: Please, use our python bug_report_model.py template.


Colab Notebook:
`* CUDA:

GPU:
available:         False
version:           None



Packages:

numpy:             1.18.5
pyTorch_debug:     False
pyTorch_version:   1.8.0a0+d0df29a
pytorch-lightning: 0.8.1
tqdm:              4.41.1



System:

OS:                Linux
architecture:

64bit



processor:         x86_64
python:            3.6.9
version:           #1 SMP Thu Jul 23 08:00:38 PDT 2020`