The models provided in the Torchvision library of PyTorch give NaN output when performing inference with CUDA on the Jetson Nano (Jetpack 4.2). Code below to reproduce:
import torch import torchvision from torchvision.models import resnet18 net = resnet18(pretrained=True).eval().cuda() input = torch.ones([1, 3, 48, 48]).cuda() with torch.no_grad(): output = net(input) print(output)
I’ve installed PyTorch and Torchvision using instructions found here: https://devtalk.nvidia.com/default/topic/1049071/jetson-nano/pytorch-for-jetson-nano/
Running the same code on the CPU gives a proper output.
Any help would be greatly appreciated!