Fixed DDP bug

tatp22 · Oct 10, 2020 · f10226a · f10226a
1 parent c3e1a6a
commit f10226a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/linformer_pytorch/linformer_pytorch.py b/linformer_pytorch/linformer_pytorch.py
@@ -161,7 +161,7 @@ def forward(self, Q, K, V, **kwargs):
                 K = self.E(K)
         Q = torch.matmul(Q, K)
 
-        P_bar = Q/torch.sqrt(torch.tensor(self.dim).type(Q.type()))
+        P_bar = Q/torch.sqrt(torch.tensor(self.dim).type(Q.type())).to(Q.device)
         if self.causal_mask is not None:
             self.causal_mask = self.causal_mask.to(Q.device)
             P_bar = P_bar.masked_fill_(~self.causal_mask, float('-inf'))