def METHOD_NAME(batch, device, non_blocking=False):
    if isinstance(batch, (list, tuple)):
        return type(batch)([METHOD_NAME(u, device, non_blocking) for u in batch])
    elif isinstance(batch, dict):
        return type(batch)([(k, METHOD_NAME(v, device, non_blocking))
                            for k, v in batch.items()])
    elif isinstance(batch, torch.Tensor):
        return batch.to(device, non_blocking=non_blocking)
    return batch