Why don’t you use the module nn.Embedding?
input_embedding = nn.Embedding(dim1,dim2)
output_embedding =  nn.Embedding(dim2,dim1)
output_embedding.weight.data = input_embedding.weight.data.transpose(0,1)

I think that way, both input and output are sharing the same storage. If you modify input_embedding, the output_embedding will be changed as well.