examples/research_projects/bertabs/modeling_bertabs.py/PositionalEncoding/forward
class PositionalEncoding:    def forward(self, emb, step=None):
        emb = emb * math.sqrt(self.dim)
        if step:
            emb = emb + self.pe[:, step][:, None, :]

        else:
            emb = emb + self.pe[:, : emb.size(1)]
        emb = self.dropout(emb)
        return emb