def concat_examples(batch, device=None, padding=0):
if len(batch) == 0:
raise ValueError('batch is empty')
if device is None:
def to_device(x):
return x
elif device < 0:
to_device = cuda.to_cpu
else:
def to_device(x):
return cuda.to_gpu(x, device, cuda.Stream.null)
first_elem = batch[0]
if isinstance(first_elem, tuple):
result = []
if not isinstance(padding, tuple):
padding = [padding] * len(first_elem)
for i in six.moves.range(len(first_elem)):
result.append(to_device(_concat_arrays(
[example[i] for example in batch], padding[i])))
return tuple(result)
评论列表
文章目录