def dict_batch_to_device(batch, device):
  """
  Moves a batch of data to device
  Args:
    - batch: Can be a Torch tensor or a dict where the values are torch tensors
    - device: A Torch device to move all the tensors to
  Returns:
    - a batch of the same type as input batch but on the device
      If a dict, also a dict with same keys
  """
  try:
    return batch.to(device)
  except AttributeError:
    # don't have a to function, must be a dict, recursively move to device
    if isinstance(batch, dict):
      return {k: dict_batch_to_device(v, device) for k, v in batch.items()}
    elif isinstance(batch, tuple):
      return tuple(dict_batch_to_device(x, device) for x in batch)
    elif isinstance(batch, list):
      return [dict_batch_to_device(x, device) for x in batch]
    else:
      raise NotImplementedError