<jupyter_start><jupyter_text>Camera  Task<jupyter_code>#import torchvision.transforms as transforms

TASK = 'object_behavior'
CATEGORIES = ['background1','redlight','greenlight','bottle']
DATASETS = ['A']

import torch
import torchvision

#device = torch.device('cuda')
output_dim = len(CATEGORIES)  #number of categories

# RESNET 18
model_class = torchvision.models.resnet18(pretrained=False)
model_class.fc = torch.nn.Linear(512, output_dim)
model_class = model_class.cuda().eval().half()
#model_class = model_class.to(device)

model_class.load_state_dict(torch.load('classification_behavior_model.pth'))



from torch2trt import torch2trt

data = torch.zeros((1, 3, 224, 224)).cuda().half()

model_trt = torch2trt(model_class, [data], fp16_mode=True)
torch.save(model_trt.state_dict(), 'trt_classification_behavior_model.pth')<jupyter_output><empty_output>