Sorry for late reply as my gmail missed Pytorch notifications.
Thanks for your help.
The  following worked
for i in range(training_iter+1):
optimizer.zero_grad()
train_x= train_x.requires_grad_(True)
output = model(train_x)
loss = -mll(output,train_y)
loss.backward(retain_graph=True)
dy_dx = torch.autograd.grad(loss.mean(), train_x,allow_unused=True, retain_graph=True)