File size: 736 Bytes
91c1e9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | """Move M4SEG batches without changing mask or index dtypes."""
import torch
def move_batch(batch, device):
def move(value):
if torch.is_tensor(value):
return value.to(device=device, non_blocking=True)
if isinstance(value, dict):
return {key: move(item) for key, item in value.items()}
if isinstance(value, list):
return [move(item) for item in value]
if isinstance(value, tuple):
return tuple(move(item) for item in value)
return value
result = move(batch)
for key in ('global_enc_images', 'grounding_enc_images'):
if result[key] is not None:
result[key] = result[key].to(dtype=torch.bfloat16)
return result
|