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