in code/pretrained_model.py [0:0]
def transform_fn(model, data, content_type, output_content_type):
transform = transforms.Compose([
transforms.Resize((128,128)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
image = np.load(BytesIO(data))
image = Image.fromarray(image)
image = transform(image)
image = image.unsqueeze(0)
#forward pass
prediction = model(image)
#get prediction
predicted_class = prediction.data.max(1, keepdim=True)[1]
response_body = np.array(predicted_class.cpu()).tolist()
return response_body, output_content_type