"""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