def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """
    Softmax, but without automatic casting to fp32 when the input is of type bfloat16
    """
    d = t.dtype
    if d is torch.bfloat16 and not is_deepspeed_initialized():
        with torch.cuda.amp.autocast(enabled=False):
            s = torch.nn.functional.softmax(t, dim=dim)
    else:
        s = torch.nn.functional.softmax(t, dim=dim)

    return s

    Softmax, but without automatic casting to fp32 when the input is of type bfloat16
    
dpkd/transformers/src/transformers/models/esm/modeling_esmfold.py