Function to_device

deepspeed-telechat/utils/utils.py:19–26

(batch, device)

def to_device(batch, device):
    output = {}
    for k, v in batch.items():
        try:
            output[k] = v.to(device)
        except:
            output[k] = v
    return output

def get_dtype_size(dtype):
    if dtype == torch.bool: