Add custom embedding (#53)

- Add custom embedding (11ef03e9e6dbfd4f118f904a57e726d99d022c14)


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)