[Python] Transferring torch tensors between devices causes data corruption?

When I transfer tensor from one cuda to another, my data is corrupted：

doc_embed_i = doc_embeds.unsqueeze(0).to(TRAIN_TYPE) # [1, D]
print(f"doc_embed_i: {doc_embed_i}")
# The printed results are as follows:
# doc_embed_i: tensor([[-1.1865e-01, 1.0254e-01, -4.8828e-02,......-2.3560e-02, 7.0801e-03, -5.2002e-02, 7.8125e-02, -8.4961e-02]],
#device='cuda:4',dtype=torch.bfloat16)

doc_embed_i = doc_embed_i.to(self.devices['translator'])
print(f"doc_embed_i: {doc_embed_i}")
# The printed results are as follows:
#doc_embed_i: tensor([[0., 1., 0., 1., 0., 1., 0., 1.,......0., 1., 0., 1., 0., 1.]],
#device='cuda:5',dtype=torch.bfloat16)


I can't find any related questions, what could be the problem?

Continue reading...