Two things:

Make sure that the input dimensions of o and mem are the same. So if you change this line:




 rectified_sigmoid:

o1 = decoder(o,mem)


to o1 = decoder(mem, mem) it would eliminate the error. This may not be what you want but shows that tgt and memory should have the same dimensions.
Also, we usually use a nn.Linear layer to fine-tune the model for our specific task. In that case, you may not need a separate decoder.






 rectified_sigmoid:

out = torch.softmax(out)


You’ll need to change this to something like out = torch.softmax(out, dim=-1)