matorus/replit-coder

1import torch2import torch.nn as nn3import torch.nn.functional as F4from torch import Tensor5 6class 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)