Reporting a minor bug

[hubert0527]

Hi, thanks for making the codes public!
I found a minor bug here. The variable self.pos_embed keeps the CPU version of the positional embedding. This is the root cause of why you need to call .to() during forward pass. To fix it, you can instead call x = x + self.pos_embed_1, which self.pos_embed_1 is the correct GPU copy auto-created by PyTorch.

This bug causes additional CPU-GPU communication time during training, but I am not quite sure how much does this costs in reality.

[yifanjiang19]

Thanks so much! I'll fix it soon