def METHOD_NAME(self, input_batch: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]]):
    """
    This function captures layer-outputs and renames them as per the AIMET exported model.
    :param input_batch: Batch of inputs for which we want to obtain layer-outputs.
    :return: layer-output name to layer-output batch dict
    """
    # Run in inference mode
    outs = self.intermediate_model(input_batch, training=False)
    output_pred = [out.numpy() for out in outs]
    return dict(zip(self.original_name_to_modified_name_mapper.values(), output_pred))