def __get_outputs_from_relay()

in src/neo_loader/abstract_model_loader.py [0:0]

• 15 lines of code
• 5 McCabe index (conditional complexity)

    def __get_outputs_from_relay(self) -> List[Dict[str, Any]]:
        import tvm
        from tvm import relay
        from tvm import ir

        mod = relay.transform.InferType()(self._relay_module_object)

        if type(mod['main'].ret_type) is ir.tensor_type.TensorType:
            dshape = []
            for item in mod['main'].ret_type.shape:
                if type(item).__name__ == 'Any':
                    dshape.append(-1)
                else:
                    dshape.append(item.value)
            return [{'dtype': mod['main'].ret_type.dtype, 'shape': dshape}]
        else:
            return [{'dtype': out.dtype, 'shape': self.__convert_relay_shape(out.shape)} for out in mod['main'].ret_type.fields]