I was kindly told that my understanding of how pytorch gradient worked was completely off the mark. Now with some help I solved my problem by using this code
net.zero_grad()
output = net(input)

g = torch.zeros(batch_size, 10, 3, 32, 32)

for i in range(10):
    g[:, i] = torch.autograd.grad(output[:, i].sum(), input, retain_graph=True)[0].data

print("GRADIENT: {}")
print(g[0])

It is probably not super helpful but maybe it will help someone in the future. I am not sure that net.zero_grad() is required now since I use grad.
Btw, is there a place where I can read how autograd works in detail, without having to read the source code?