Add custom embedding (#22)

- Add custom embedding (f0249f31962c2e1dbba35c1ad7a8f57efceb19df)


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)