I came across the same problem while trying to implement something where I wanted to use the embedding as part of the loss function. I ended up implementing a new Embedding class that works by one hot encoding the inputs with concept similar to resonance amplitude calculation and then performing a matrix multiplication. My embedding looks like this
class HotEmbedding(torch.nn.Module):
    def __init__(self, max_val, embedding_dim, eps=1e-2):
        super(HotEmbedding, self).__init__()
        self.A = torch.arange(max_val, requires_grad=False)
        self.B = torch.randn((max_val, embedding_dim), requires_grad=True)
        self.eps = eps

    def forward(self, x):
        return 1/((x.unsqueeze(1)**2 - self.A**2)+self.eps) @ self.B

And you can use it as
layer = HotEmbedding(10, 5)
x = torch.tensor([1.,2.,3.,1.,2.,3.], requires_grad=True)
y = layer(x)

x and y look like this
tensor([1., 2., 3., 1., 2., 3.], requires_grad=True)
tensor([[  16.7347, -117.7052,   10.8307,   14.4596,  -14.4507],
        [  55.2739,    8.8376,  -12.5775, -121.5673, -114.9303],
        [ -86.5265,  -92.7613,   84.4617,   68.5347,  203.2270],
        [  16.7347, -117.7052,   10.8307,   14.4596,  -14.4507],
        [  55.2739,    8.8376,  -12.5775, -121.5673, -114.9303],
        [ -86.5265,  -92.7613,   84.4617,   68.5347,  203.2270]],
       grad_fn=<MmBackward>)