    def forward(self, x, lengths=None):
        """Average pooling across time steps (dim=1) with optionally lengths.
        Args:
            x: torch.Tensor of shape (N, T, ...)
            lengths: None or torch.Tensor of shape (N,)
            dim: dimension to pool
        """
        if lengths is None:
            return x.mean(dim=1, keepdim=False)
        else:
            mask = get_mask_from_lengths(lengths).type(x.type()).to(x.device)
            mask_shape = list(mask.size()) + [1 for _ in range(x.ndimension()-2)]
            mask = mask.reshape(*mask_shape)
            numer = (x * mask).sum(dim=1, keepdim=False)
            denom = mask.sum(dim=1, keepdim=False)
            return numer / denom
Average pooling across time steps (dim=1) with optionally lengths.
Args:
x: torch.Tensor of shape (N, T, ...)
lengths: None or torch.Tensor of shape (N,)
dim: dimension to pool
examples/textless_nlp/gslm/unit2speech/tacotron2/layers.py