I think your way to permute the dimensions might not work, if you provide a single sample with a batch size of 1.
Currently you are squeezing and unsqueezing the output of your embedding to change the dimensions to [seqlen, batch, input].
However, if your batch size is only 1, you’ll squeeze this dimension also.
A better approach would be to use out = out.permute(1, 0, 2) and call contiguous() on it if necessary.
Could you try that and see, if it’s working?