def dropout_condition(prob: float, txt: torch.Tensor, null_txt: torch.Tensor) -> torch.Tensor:
    """
    Apply dropout to the text tensor.

    Args:
        prob (float): The dropout probability.
        txt (torch.Tensor): The text tensor.
        null_txt (torch.Tensor): The null text tensor.

    Returns:
        torch.Tensor: The text tensor with dropout applied.
    """
    if prob == 0:
        warnings.warn("Dropout probability is 0, skipping dropout")
    drop_ids = torch.rand(txt.shape[0], device=txt.device) < prob
    drop_ids = drop_ids.view((drop_ids.shape[0],) + (1,) * (txt.ndim - 1))
    new_txt = torch.where(drop_ids, null_txt, txt)
    return new_txt

    Apply dropout to the text tensor.

    Args:
        prob (float): The dropout probability.
        txt (torch.Tensor): The text tensor.
        null_txt (torch.Tensor): The null text tensor.

    Returns:
        torch.Tensor: The text tensor with dropout applied.
    
opensora/utils/train.py