code: def shake_shake_branch(x, output_filters, stride, rand_forward, rand_backward,
                         hparams):
    
    is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
    x = tf.nn.relu(x)
    x = tf.layers.conv2d(
        x,
        output_filters, (3, 3),
        strides=(stride, stride),
        padding="SAME",
        name="conv1")
    x = tf.layers.batch_normalization(x, training=is_training, name="bn1")
    x = tf.nn.relu(x)
    x = tf.layers.conv2d(x, output_filters, (3, 3), padding="SAME", name="conv2")
    x = tf.layers.batch_normalization(x, training=is_training, name="bn2")
    if is_training:
      x = x * rand_backward + tf.stop_gradient(x * rand_forward -
                                               x * rand_backward)
    else:
      x *= 1.0 / hparams.shake_shake_num_branches
    return x
docs: Building a 2 branching convnet.