def model_inference_fn(features, training, params):
    """Builds just the inference part of the model graph.

    Args:
        features: input features tensor.
        training: True if the model is training.
        params: A dictionary

    Returns:
        (policy_output, value_output, logits) tuple of tensors.
    """

    mg_batchn = functools.partial(
        tf.layers.batch_normalization,
        axis=-1,
        momentum=.95,
        epsilon=1e-5,
        center=True,
        scale=True,
        fused=True,
        training=training)

    mg_conv2d = functools.partial(
        tf.layers.conv2d,
        filters=params['conv_width'],
        kernel_size=3,
        padding="same",
        data_format="channels_last",
        use_bias=False)

    mg_global_avgpool2d = functools.partial(
        tf.layers.average_pooling2d,
        pool_size=go.N,
        strides=1,
        padding="valid",
        data_format="channels_last")

    def mg_activation(inputs):
        if FLAGS.use_swish:
            return tf.nn.swish(inputs)

        return tf.nn.relu(inputs)


    def residual_inner(inputs):
        conv_layer1 = mg_batchn(mg_conv2d(inputs))
        initial_output = mg_activation(conv_layer1)
        conv_layer2 = mg_batchn(mg_conv2d(initial_output))
        return conv_layer2

    def mg_res_layer(inputs):
        residual = residual_inner(inputs)
        output = mg_activation(inputs + residual)
        return output

    def mg_squeeze_excitation_layer(inputs):
        # Hu, J., Shen, L., & Sun, G. (2018). Squeeze-and-Excitation Networks.
        # 2018 IEEE/CVF Conference on Computer Vision, 7132-7141.
        # arXiv:1709.01507 [cs.CV]

        channels = params['conv_width']
        ratio = FLAGS.SE_ratio
        assert channels % ratio == 0

        residual = residual_inner(inputs)
        pool = mg_global_avgpool2d(residual)
        fc1 = tf.layers.dense(pool, units=channels // ratio)
        squeeze = mg_activation(fc1)

        if FLAGS.use_SE_bias:
            fc2 = tf.layers.dense(squeeze, units=2*channels)
            # Channels_last so axis = 3 = -1
            gamma, bias = tf.split(fc2, 2, axis=3)
        else:
            gamma = tf.layers.dense(squeeze, units=channels)
            bias = 0

        sig = tf.nn.sigmoid(gamma)
        # Explicitly signal the broadcast.
        scale = tf.reshape(sig, [-1, 1, 1, channels])

        excitation = tf.multiply(scale, residual) + bias
        return mg_activation(inputs + excitation)

    initial_block = mg_activation(mg_batchn(mg_conv2d(features)))

    # the shared stack
    shared_output = initial_block
    for _ in range(params['trunk_layers']):
        if FLAGS.use_SE or FLAGS.use_SE_bias:
            shared_output = mg_squeeze_excitation_layer(shared_output)
        else:
            shared_output = mg_res_layer(shared_output)

    # Policy head
    policy_conv = mg_conv2d(
        shared_output, filters=params['policy_conv_width'], kernel_size=1)
    policy_conv = mg_activation(mg_batchn(policy_conv, center=False, scale=False))
    logits = tf.layers.dense(
        tf.reshape(
            policy_conv, [-1, params['policy_conv_width'] * go.N * go.N]),
        go.N * go.N + 1)

    policy_output = tf.nn.softmax(logits, name='policy_output')

    # Value head
    value_conv = mg_conv2d(
        shared_output, filters=params['value_conv_width'], kernel_size=1)
    value_conv = mg_activation(mg_batchn(value_conv, center=False, scale=False))

    value_fc_hidden = mg_activation(tf.layers.dense(
        tf.reshape(value_conv, [-1, params['value_conv_width'] * go.N * go.N]),
        params['fc_width']))
    value_output = tf.nn.tanh(
        tf.reshape(tf.layers.dense(value_fc_hidden, 1), [-1]),
        name='value_output')

    return policy_output, value_output, logits