defined input_fn in inference.py but not getting used?

[kvnsng]

Hello,
I'm trying to deploy a PyTorch Model to SageMaker Endpoint, but the documentation is very confusing. I've defined a new inference.py with the following input_fn and output_fn:

def input_fn(request_body, request_content_type):
    if request_content_type == 'application/json':
        data = json.loads(request_body)
        data = np.array(data['Input'])
        if data.shape[1] == 3:
            return torch.Tensor(data)
        elif data.shape[3] == 3:
            return torch.Tensor(data).permute((0,3,1,2))
        else:
            raise Exception('image must have 3 channels [3, width, height]')
    else:
        # Handle other content-types here or raise an Exception
        # if the content type is not supported.
        raise Exception("Unsupported content type: " + request_content_type)

def output_fn(prediction_output, content_type):
    pred = non_max_suppression(prediction_output[0], 0.5, 0.45, None, False, 1000)
    return json.dumps(pred[0].tolist())

I've deployed the model locally using the following commands:

import os
from sagemaker.pytorch import PyTorchModel
from sagemaker.local import LocalSession

sagemaker_session = LocalSession()
sagemaker_session.config = {'local': {'local_code': True}}

model = PyTorchModel(entry_point='inference.py',
                     model_data=model_data, 
                     framework_version='1.11', 
                     py_version='py38',
                     role=role)

INSTANCE_TYPE = 'local'

predictor = model.deploy(initial_instance_count=1, 
                         instance_type=INSTANCE_TYPE,)

but when I try to call the model by using predictor.predict(data), I get the following error:

---------------------------------------------------------------------------
UnpicklingError                           Traceback (most recent call last)
~/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/numpy/lib/npyio.py in load(file, mmap_mode, allow_pickle, fix_imports, encoding)
    459             try:
--> 460                 return pickle.load(fid, **pickle_kwargs)
    461             except Exception:

UnpicklingError: unpickling stack underflow

During handling of the above exception, another exception occurred:

OSError                                   Traceback (most recent call last)
/tmp/ipykernel_15312/4230898929.py in <module>
----> 1 predictor.predict(io.BytesIO(data))

~/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/sagemaker/predictor.py in predict(self, data, initial_args, target_model, target_variant, inference_id)
    160         )
    161         response = self.sagemaker_session.sagemaker_runtime_client.invoke_endpoint(**request_args)
--> 162         return self._handle_response(response)
    163 
    164     def _handle_response(self, response):

~/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/sagemaker/predictor.py in _handle_response(self, response)
    166         response_body = response["Body"]
    167         content_type = response.get("ContentType", "application/octet-stream")
--> 168         return self.deserializer.deserialize(response_body, content_type)
    169 
    170     def _create_request_args(

~/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/sagemaker/deserializers.py in deserialize(self, stream, content_type)
    223                 return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
    224             if content_type == "application/x-npy":
--> 225                 return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
    226         finally:
    227             stream.close()

~/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/numpy/lib/npyio.py in load(file, mmap_mode, allow_pickle, fix_imports, encoding)
    460                 return pickle.load(fid, **pickle_kwargs)
    461             except Exception:
--> 462                 raise IOError(
    463                     "Failed to interpret file %s as a pickle" % repr(file))
    464     finally:

OSError: Failed to interpret file <_io.BytesIO object at 0x7f7c90859090> as a pickle

I've tried changing the data to be a numpy array, json, byte buffer and none of them seem to work. It seems like the default serializer is being used instead of the input_fn in my inference.py?

Any suggestions/help will be much appreciated. Thank you!