def get_channel_dim(input_tensor, data_format='INVALID'):
  """Returns the number of channels in the input tensor."""
  shape = input_tensor.get_shape().as_list()
  assert data_format != 'INVALID'
  assert len(shape) == 4
  if data_format == 'NHWC':
    return int(shape[3])
  elif data_format == 'NCHW':
    return int(shape[1])
  else:
    raise ValueError('Not a valid data_format', data_format)