from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image
from transformers import AutoTokenizer

model_name = "DenisNovac/nsfw_image_detection"

model = AutoModelForImageClassification.from_pretrained(model_name, torchscript=True, return_dict=False)

processor = AutoImageProcessor.from_pretrained(model_name)

image = Image.open("images/hentai.jpg")
image_inputs = processor(images=image, return_tensors="pt")

config = {'forward': [image_inputs['pixel_values']]}
converted = torch.jit.trace_module(model,  config)

torch.jit.save(converted, "converted-to-torchscript.pt")