import torch
torch.manual_seed(0)
dev = torch.device('cuda')
embedding_layer = torch.nn.Embedding(10, 5, padding_idx=0)
# random_ix = torch.randint(high=10, size=(5,))
# random_ix = random_ix.to(dev)
random_ix = torch.LongTensor([8, 8, 8, 4, 8, 7]).to(dev)
print(random_ix)
def hook(module, grad_in, grad_out):
    print("GRAD OUT")
    print(grad_out[0])
    print(grad_out[0].grad_fn)
    print("GRAD IN")
    print(grad_in[0][random_ix])
embedding_layer.register_backward_hook(hook)
embedding_layer.to(dev)
embeds = embedding_layer(random_ix)
print(embeds)
print(embedding_layer.weight)
# embeds.register_hook(lambda grads: print('hook', grads))
merged = torch.sum(embeds, dim=1)
summed = merged.sum()
grad_auto = torch.autograd.grad(summed, embedding_layer.weight, create_graph=True)
summed.backward()
print("EMBEDDING LAYER WEIGHT GRADS")
print(embedding_layer.weight.grad)
print("GRAD AUTO GRADS")
print(grad_auto[0][random_ix])