src/transformers/models/esm/modeling_esmfold.py/EsmFoldDropout/__init__
class EsmFoldDropout:    def __init__(self, r: float, batch_dim: Union[int, List[int]]):
        super().__init__()

        self.r = r
        if type(batch_dim) == int:
            batch_dim = [batch_dim]
        self.batch_dim = batch_dim
        self.dropout = nn.Dropout(self.r)