sea-lion-3b / custom_embedding.py
xianbin's picture
Add 3B model files
430b3b1
raw
history blame contribute delete
292 Bytes
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)