model_wrapper

torch_to_nnef.model_wrapper

Wrap model to bypass limitation of torch_to_nnef internals.

Cases where inputs or outputs of a model contains:

tuples, list, dicts, Object.

UnfoldModelInfo dataclass

UnfoldModelInfo(model: nn.Module, original_inputs: T.Tuple[torch.Tensor], original_outputs: T.List[torch.Tensor], flat_inputs: T.Tuple[torch.Tensor], flat_outputs: T.Tuple[torch.Tensor], input_names: T.List[str], output_names: T.List[str])

Hold model input/output structure information.

WrapStructIO

WrapStructIO(model: nn.Module, input_infos, output_infos)

Bases: Module

Once traced it should be nop in final graph.