Base Export Module.

params A dataclass for parameters to the module.
model A model instance which contains weights and forward computation.
inference_step An optional callable to forward-pass the model. If not specified, it creates a parital function with model as an required kwarg.
preprocessor An optional callable to preprocess the inputs.
postprocessor An optional callable to postprocess the model outputs.



View source

Get defined function signatures.


View source

The bare inference function which should run on all devices.

Expecting tensors are passed in through keyword arguments. Returns a dictionary of tensors, when the keys will be used inside the SignatureDef.