Create custom_embedding.py (#37)

- Create custom_embedding.py (4c8fc62bd3bee3527654c502f211795e7757e3ea)


Co-authored-by: Brandon Cui <bcui19@users.noreply.huggingface.co>

custom_embedding.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class SharedEmbedding(nn.Embedding):
+
+    def forward(self, input: Tensor, unembed: bool = False) -> Tensor:
+        if unembed:
+            return F.linear(input, self.weight)
+        return super().forward(input)