import torch


def prepare_datas_and_labels_mgpu(data_iterator, gpu_ids):
  """Prepare datas and labels for multi-gpu computation.

  Args:
    data_iterator: An Iterator instance of pytorch.DataLoader, which
      return a dictionary of `datas`, `labels`, and a scalar of `index`.
    gpu_ids: A list of scalars indicates the GPU device ids.

  Return:
    A list of tuples of `datas` and `labels`.
  """
  input_batch, label_batch = [], []
  for gpu_id in gpu_ids:
    data, label, index = data_iterator.next()
    for k, v in data.items():
      data[k] = (v if not torch.is_tensor(v)
                  else v.pin_memory().to(gpu_id, non_blocking=True))
    for k, v in label.items():
      label[k] = (v if not torch.is_tensor(v)
                   else v.pin_memory().to(gpu_id, non_blocking=True))
    input_batch.append(data)
    label_batch.append(label)

  return input_batch, label_batch