Orin and compatibility with torchvision models

Hello, currently I’m developing using a torchvision network on the Orin, running my model within the docker image


Which is found here


The issue I’m currently facing is that I am able to successfully run data through torchvision networks, however the output is invalid/null/nan when running on the orin, however the output is valid when running on my x86_64 machine.

Here is example code that is able to reproduce this, keep in mind I have tested this across other models in the torchvision network as well, such as

#this keypoint model returns no keypoints when ran on the orin, where the same input data returns keypoints on an x86_64 machine

Example code to reproduce the issue

import torchvision
import torchvision.transforms as transforms

# Load the pre-trained model
model = torchvision.models.resnet50(pretrained=True)

# Define the transform
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

# Load the CIFAR-10 dataset
dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# Use the first image from the dataset
image, label = dataset[23]

print("This is a picture of a ", label, "in CIFAR-10 labels. It is a fire truck")

# Add a batch dimension
image = image.unsqueeze(0)

# Run the image through the model
prediction = model(image).squeeze(0).softmax(0)
class_id = prediction.argmax().item()
score = prediction[class_id].item()
print(f"{class_id}: {100 * score:.1f}%")
print("This is a picture of a ", class_id, "in IMGNET labels. 555 is a firetruck...")

@samanthaxpalmer does it work if you run the model on the GPU? You can find the code for a similar test I run on torchvision in the l4t-pytorch container here:

It runs Alexnet, Googlenet, ResNet18, and ResNet50 on 5000 test images and confirms that the inferencing accuracy is what PyTorch publishes for those models. (These tests all passed prior to the container being released)

Thank you so much, I knew it had to be something very simple

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.