def _skip_first_max_pooling(inputs):
  """Whether to skip the first max pooling layer in ResNet.

  For 128x128 inputs, we skip the first 3x3 2-stride max pooling layer.

  Args:
    inputs: The inputs passed to ResNet.

  Returns:
    Whether to skip the first max pooling layer in ResNet.
  """
  dims = inputs.shape.dims
  size = dims[1]  # dims is like [B, H, W, C]
  return size == 128