Code to reproduce error on https://github.com/aws/aws-neuron-sdk/issues/369

compile_fbnet.py

import torch
import timm

import torch.neuron

# original model
model = timm.create_model("fbnetc_100", pretrained=True)
model.eval()

# number of parameters
pytorch_total_params = sum(p.numel() for p in model.parameters())
print("total_param: " + str(pytorch_total_params))

# dummy data
data_in = torch.rand(1, 3, 512, 512)
print(model(data_in))

# trace the model
print("Trace model")
torch.neuron.analyze_model(model, example_inputs=[data_in])
model_neuron = torch.neuron.trace(model, example_inputs=[data_in])
model_neuron.save("test.pt")

# load the compiled model and run inference
print("Run inference on neuron")
model_neuron_loaded = torch.jit.load("test.pt")
print(model_neuron_loaded(data_in))