I am using a boolean tensor to index another tensor. I am building a very basic Deep Q Learning model with experience replay, the below code is meant to NOT the mask denoting which state transitions are terminal, using this to index the state value estimates tensor and place the needed values there.
In the end the tensor should be arranged such that if terminal_mask[i]==1 then discounted_Q-estimates = reward otherwise discounted_Q_estimates = reward + GAMMA* max_a’(Q(s,a’)).
The error this returns is:
Traceback (most recent call last):
  File "C:\Users\Arham\Desktop\DeepQBreakout\main.py", line 121, in <module>
    discounted_Q_estimates[torch.logical_not(terminal_mask)] = GAMMA * ns_preds
RuntimeError: shape mismatch: value tensor of shape [32] cannot be broadcast to indexing result of shape [31]

which is strange because every print statement in the below snippet returns “torch.size([32])” indicating that the value and index tensors have the same size, can anyone help me debug this? Thank you!
states, actions, rewards, next_states, terminal_mask = replay_buffer.sample()
          #if terminal state label is reward else reward + discounted Q estimate
          print(terminal_mask.shape) #torch.Size([32])
          print(torch.logical_not(terminal_mask).shape) #torch.Size([32])
          with torch.no_grad():
              ns_preds, _ = torch.max(model(next_states.to(device)), dim=1)
              discounted_Q_estimates = torch.zeros(BATCHSIZE, device=device)
              print(ns_preds.shape) #torch.Size([32])
              print(discounted_Q_estimates.shape) #torch.Size([32])
              discounted_Q_estimates[torch.logical_not(terminal_mask)] = GAMMA * ns_preds
              y = rewards + discounted_Q_estimates