docs: def per_device_batch_size(batch_size, num_gpus):
    """For multi-gpu, batch-size must be a multiple of the number of GPUs.
  
    Note that this should eventually be handled by DistributionStrategies
    directly. Multi-GPU support is currently experimental, however,
    so doing the work here until that feature is in place.
  
    Args:
      batch_size: Global batch size to be divided among devices. This should be
        equal to num_gpus times the single-GPU batch_size for multi-gpu training.
      num_gpus: How many GPUs are used with DistributionStrategies.
  
    Returns:
      Batch size per device.
  
    Raises:
      ValueError: if batch_size is not divisible by number of devices
    """
    if num_gpus <= 1:
      return batch_size
  
    remainder = batch_size % num_gpus
    if remainder:
      err = ('When running with multiple GPUs, batch size '
             'must be a multiple of the number of available GPUs. Found {} '
             'GPUs with a batch size of {}; try --batch_size={} instead.'
            ).format(num_gpus, batch_size, batch_size - remainder)
      raise ValueError(err)
    return int(batch_size / num_gpus)