Add custom embedding (#42)

- Add custom embedding (216226e780056807fd082a6e5f41a4ece987b49c)


Co-authored-by: Brandon Cui <[email protected]>

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)