LookupTable in PyTorch

It seems Embedding is the way to go since they have the same input and output.

However, I am having trouble translating this line:

weight_lut.data.module.updateGradInput = function(self, input) return input end

to PyTorch.

Do I need to subclass Emdedding an supply a custom backward function?