replit/replit-code-v1_5-3b

1import torch.nn as nn2import torch.nn.functional as F3from torch import Tensor4 5class SharedEmbedding(nn.Embedding):6 7    def forward(self, input: Tensor, unembed: bool=False) -> Tensor:8        if unembed:9            return F.linear(input, self.weight)10        return super().forward(input)