File "<ipython-input-1-68fc3e10bc2f>", line 141, in update
    action = brain.update(last_reward, last_signal) # playing the action from our ai (the object brain of the dqn class)

  File "/media/ozzycodes/ExtraData/Documents/Programming Files/Artificial Intelligence/OzzyCodes-AI/ozzyAI.py", line 100, in update
    self.learn(batch_state, batch_next_state, batch_reward, batch_action)

  File "/media/ozzycodes/ExtraData/Documents/Programming Files/Artificial Intelligence/OzzyCodes-AI/ozzyAI.py", line 83, in learn
    outputs = self.model(batch_state).gather(1, batch_action.unsqueeze(1)).squeeze(1) # Gathers the best action to play

RuntimeError: Expected object of scalar type Long but got scalar type Float for argument #3 'index'

def learn(self, batch_state, batch_next_state, batch_reward, batch_action):
    outputs = self.model(batch_state).gather(1, batch_action.unsqueeze(1)).squeeze(1) # RuntimeError
    next_outputs = self.model(batch_next_state).detach().max(1)[0]
    target = self.gamma * next_outputs + batch_reward
    td_loss = F.smooth_l1_loss(outputs, target) # td = temporal difference
    self.optimizer.zero_grad()
    td_loss.backward(retain_variables = True) # Free some memory when backtracking
    self.optimizer.step()