If your embedding is standard then you can use torch.matmul as seen here https://github.com/facebookresearch/EGG/blob/170e5fe63c13244121a5b29b9bfb4870a0f11796/egg/core/gs_wrappers.py#L203
Otherwise, you have to also implement the Embedding. E.g. if your padding_idx != 0, I’m not sure but I think you have to figure out how to make it not contribute to the gradient