Solved by telecom in post #2 
              
              
                If someone intersted I resolved this by doing like that: 
def threeto2ddoc(doc):
    doc = torch.cat((doc, torch.zeros(doc.shape[0],doc.shape[1],doc.shape[0]-1)),dim=2)
    for i in range(doc.shape[0]):
        for j in range(doc.shape[1]):
            doc[i,j,:]=torch.cat((doc[i,j,-i:],doc[i,j,0:-i…