def transform_fn()

in code/pretrained_model_with_debugger_hook.py [0:0]


def transform_fn(model_with_hook, data, content_type, output_content_type):
    model = model_with_hook.model
    hook = model_with_hook.hook

    val_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 = val_transform(image)

    image = image.unsqueeze(0)
    image = image.to('cpu').requires_grad_()
    if hook is not None:
        hook.image_gradients(image)

    #forward pass
    prediction = model(image)

    #get prediction
    predicted_class = prediction.data.max(1, keepdim=True)[1]
    output = prediction[0, predicted_class[0]]
    model.zero_grad()

    #compute gradients with respect to outputs 
    output.backward()

    response_body = np.array(predicted_class.cpu()).tolist()
    return response_body, output_content_type