Deep learning 维数上的移位张量

假设我有一个形状为[batch_size X length X 1024]的张量 我想做以下工作： 对于批处理中的i元素，我希望将所有嵌入的“长度”元素的（嵌入1024个）按其位置移动。 例如，向量A[0，0，：]应该保持不变，A[0，1，：]应该移动（或滚动）1，A[0，15，：]应该移动15。 这适用于批处理中的所有元素。 到目前为止，我使用for循环进行了此操作，但效率不高
下面是我的for循环代码：

    x = # [batchsize , length , 1024]
    new_embedding = []
    llist = []
    batch_size = x.shape[0]
    seq_len    = x.shape[1]
    for sample in range(batch_size):
     for token in range(seq_len):
      orig = x[sample , token , : ]
      new_embedding.append(torch.roll(orig , token , 0))
     llist.append(torch.stack(new_embedding , 0))
     new_embedding = []
    x = torch.stack(llist , 0)