fix resize_token_embeddings (#11572)

7c622482 · Stas Bekman · GitHub · fe82b1bf · 7c622482
Unverified Commit 7c622482 authored May 03, 2021 by Stas Bekman Committed by GitHub May 03, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

src/transformers/modeling_utils.py src/transformers/modeling_utils.py +3 -1

No files found.
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -682,7 +682,9 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
            )

        # Build new embeddings
-        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim).to(self.device)
+        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim).to(
+            self.device, dtype=old_embeddings.weight.dtype
+        )

        # initialize all new embeddings (in particular added tokens)
        self._init_weights(new_embeddings)