I’ve got the following mechanism for adding positional awareness to representations, it works well but  is very slow.
The fundamental issue is that each sample in the batch has a different real length and I therefore have to loop over and create different embeddings for each of them.
class PositionSymmetricLayer(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.embedding = nn.Parameter(torch.linspace(0., 1., dim//2), requires_grad=False)
        self.dim = dim

    def forward(self, x, durs):
        N = x.size(0)
        T = x.size(1)
        pos_emb = torch.zeros((N, T, self.dim), dtype=torch.float32).cuda()
        for i, dur in enumerate(durs):
            factor = torch.linspace(0, 1, dur).unsqueeze(1).cuda()
            pos = (torch.cat((self.embedding - factor, -self.embedding + factor,), dim=-1) >= 0).float()
            pos_emb[i, -dur:] = pos
        x = torch.cat((x, pos_emb,), dim=-1)
        return x

To help understand what this does you can see in the following what pos is given a dur of 3 and a dim of 8:
tensor([[1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 1., 1., 1., 1., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 1.]])

Is there something I can do to speed the for loop up?