def pad(x: np.ndarray, num_devices: int, constant: float = 0.0):
  """Pads x of shape [bs, t, h, w] into [device_bs, #device, t, h, w]."""
  assert (
      x.ndim == 4 or x.ndim == 5
  ), 'Not in [bs, t, h, w] or [bs, t, h, w, rvq]'
  num_batches = int(np.ceil(x.shape[0] * 1.0 / num_devices))
  num_to_pad = num_batches * num_devices - x.shape[0]
  pad_values = np.ones_like(x[0]) * constant
  if num_to_pad > 0:
    logging.log_first_n(logging.WARNING, 'Padding %d frames with shape %s.', 20,
                        num_to_pad, pad_values.shape)
    padded_x = np.concatenate((x, [pad_values] * num_to_pad))
  else:
    padded_x = x
  padded_x = padded_x.reshape([-1, num_devices, *x.shape[1:]])
  return padded_x.astype(x.dtype)
Pads x of shape [bs, t, h, w] into [device_bs, #device, t, h, w].
videogvt/interfaces/colab_utils.py