Function _maybe_cast_to_bf16

tests/models/testing_utils/utils.py:13–22

Cast model and floating-point inputs to bfloat16 when the backend requires it.

(backend, model, inputs_dict)

def _maybe_cast_to_bf16(backend, model, inputs_dict):
    """Cast model and floating-point inputs to bfloat16 when the backend requires it."""
    if not backend or backend not in _BF16_REQUIRED_BACKENDS:
        return model, inputs_dict
    model = model.to(dtype=torch.bfloat16)
    inputs_dict = {
        k: v.to(dtype=torch.bfloat16) if isinstance(v, torch.Tensor) and v.is_floating_point() else v
        for k, v in inputs_dict.items()
    }
    return model, inputs_dict