src/transformers/models/esm/modeling_esmfold.py/EsmFoldDropout/forward
class EsmFoldDropout:    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shape = list(x.shape)
        if self.batch_dim is not None:
            for bd in self.batch_dim:
                shape[bd] = 1
        return x * self.dropout(x.new_ones(shape))