import torch


def batch_to_device(batch, device):
  """
  Transfers a tensor, list of tensors, or dict of tensors to a given device.
  Works best with pinned memory. Non-tensors are returned as they are.

  :param batch: list, dict, or single tensor
  :param device: Target device
  :return:
  """
  if isinstance(batch, torch.Tensor):
    return batch.to(device, non_blocking=True)
  elif isinstance(batch, list):
    return [batch_to_device(b, device) for b in batch]
  elif isinstance(batch, dict):
    return {k: batch_to_device(v, device) for k, v in batch.items()}
  else:
    return batch