import torch
from torch.distributed.fsdp import MixedPrecision

bfSixteen = MixedPrecision(
    param_dtype=torch.bfloat16,
    # Gradient communication precision.
    reduce_dtype=torch.bfloat16,
    # Buffer precision.
    buffer_dtype=torch.bfloat16,
    # the casting here is to make FSDP work in eval mode
    # there is a bug where model inputs aren't cast to fp32 for model.eval
    # these should be changed with EXTREME care.
    # https://github.com/pytorch/pytorch/pull/99546#discussion_r1180898687
    cast_forward_inputs=True,
    cast_root_forward_inputs=False,
)

float_32 = MixedPrecision(
    param_dtype=torch.float32,
    # Gradient communication precision.
    reduce_dtype=torch.float32,
    # Buffer precision.
    buffer_dtype=torch.float32,
)

no_policy = None