def cast_batch_to_half(batch):
    """
    Cast the float32 tensors in a batch to float16.
    It should be called before feeding the batch to the FP16 DeepSpeed model.

    Args:
        batch (torch.tensor or container of torch.tensor): input batch
    Returns:
        return_batch: same type as the input batch with internal float32 tensors casted to float16
    """
    if torch.is_tensor(batch):
        if torch.is_floating_point(batch):
            return_batch = batch.to(torch.float16)
        else:
            return_batch = batch
    elif isinstance(batch, list):
        return_batch = [cast_batch_to_half(t) for t in batch]
    elif isinstance(batch, tuple):
        return_batch = tuple(cast_batch_to_half(t) for t in batch)
    elif isinstance(batch, dict):
        return_batch = {}
        for k in batch:
            return_batch[k] = cast_batch_to_half(batch[k])
    else:
        logger.debug(f"Can not cast type {type(batch)} to float16. Skipping it in the batch.")
        return_batch = batch

    return return_batch

    Cast the float32 tensors in a batch to float16.
    It should be called before feeding the batch to the FP16 DeepSpeed model.

    Args:
        batch (torch.tensor or container of torch.tensor): input batch
    Returns:
        return_batch: same type as the input batch with internal float32 tensors casted to float16
    
trainer/utils/misc.py