def transform_fn()

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