Ontocord.AI commited on
Commit
c5e4b95
1 Parent(s): 4555156

Upload custom_embedding.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. custom_embedding.py +11 -0
custom_embedding.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+ from torch import Tensor
5
+
6
+ class SharedEmbedding(nn.Embedding):
7
+
8
+ def forward(self, input: Tensor, unembed: bool=False) -> Tensor:
9
+ if unembed:
10
+ return F.linear(input, self.weight)
11
+ return super().forward(input)