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.
Hi, thanks for making the codes public!
I found a minor bug here. The variable
self.pos_embedkeeps 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 callx = x + self.pos_embed_1, whichself.pos_embed_1is 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.