In my neural network (RNN), I am defining the loss function such that the output of the neural network is used to find the index (binary) and then the index is used to extract the required element from an array which in turn will be used to calculate MSELoss.
However, the program gives parameter().grad = None error which is mostly because the graph is breaking somewhere. What is the problem with the error function defined.
Framework: Pytorch
The codes are as follow:
Neural Network:
class RNN(nn.Module):
  def __init__(self):
    super(RNN, self).__init__()
    self.hidden_size = 8
    # self.input_size = 2
    self.h2o = nn.Linear(self.hidden_size, 1)
    self.h2h = nn.Linear(self.hidden_size, self.hidden_size)
    self.sigmoid = nn.Sigmoid()
  def forward(self,hidden):
    output = self.h2o(hidden)
    output = self.sigmoid(output)
    hidden = self.h2h(hidden)
    return output, hidden
  def init_hidden(self):
    return torch.zeros(1, self.hidden_size)

Loss Function, train step and training
rnn = RNN()
criterion = nn.MSELoss()

def loss_function(previous, output, index):
  code = 2*(output > 0.5).long()
  current = Q_m2[code:code+2, i]
  return criterion(current, previous), current

def train_step():
  hidden = rnn.init_hidden()
  rnn.zero_grad()
  # Q_m2.requires_grad = True
  # Q_m2.create_graph = True 
  loss = 0
  previous = Q_m[0:2, 0]
  for i in range(1, samples):
    output, hidden = rnn(hidden)
    l, previous = loss_function(previous, output, i)
    loss+=l
  loss.backward()
  # Q_m2.retain_grad()
  for p in rnn.parameters():
    p.data.add_(p.grad.data, alpha=-0.05)
  return output, loss.item()/(samples - 1)

def training(epochs):
  running_loss = 0
  for i in range(epochs):
    output, loss = train_step()
    print(f'Epoch Number: {i+1}, Loss: {loss}')
    running_loss +=loss

Q_m2
Q_m = np.zeros((4, samples))
for i in range(samples):
  Q_m[:,i] = q_x(U_m[:,i])
Q_m = torch.FloatTensor(Q_m)
Q_m2 = Q_m
Q_m2.requires_grad = True
Q_m2.create_graph = True

Error:
<ipython-input-36-feefd257c97a> in train_step()
     21   # Q_m2.retain_grad()
     22   for p in rnn.parameters():
---> 23     p.data.add_(p.grad.data, alpha=-0.05)
     24   return output, loss.item()/(samples - 1)
     25 

AttributeError: 'NoneType' object has no attribute 'data'