Build the ResNet model as normal up to the following lines:
inputs = tf.reshape(
inputs, [-1, 2048 if block_fn is bottleneck_block else 512])