external_embedding device to embedding device

2024-06-19 19:23:58 +03:00
parent 63ebb03725
commit 50a473c1b7
1 changed files with 1 additions and 1 deletions
@@ -375,7 +375,7 @@ class EmbeddingLayerWithFixes(nn.Module):
        start = external_embedding["start"]
        end = external_embedding["end"]
        target_ids_to_replace = [i for i in range(start, end)]
-        ext_emb = external_embedding["embedding"]
+        ext_emb = external_embedding["embedding"].to(embedding.device)

        # do not need to replace
        if not (input_ids == start).any():