PyTorch shape_mismatch error between tensors of same size?

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

Based on the error message the indexed tensor has a shape of [31], while the value tensor has a shape of [32]. This code snippet illustrates the issue:

x = torch.randn(32)
idx = torch.ones(32).bool()
x[idx] = torch.randn(32) # works

idx[0] = False
x[idx] = torch.randn(32) # fails
# > RuntimeError: shape mismatch: value tensor of shape [32] cannot be broadcast to indexing result of shape [31]

so you won’t be able to assign 32 values to the x[idx] which is indexing 31 values only.

Thank you, I figured that this was the problem but based on my print debugging it would seem that the tensors were all of size [32], what could account for this sort of error?

The shapes are also [32] in my example, but the number of True values in idx is 31.
Indexing x with 31 True values will return a tensor of the shape [31]:

x = torch.randn(32)
idx = torch.ones(32).bool()
idx[0] = False
print(x[idx].shape)
# > torch.Size([31])

which will then create the error since you are trying to assing 32 values to it.