how can i make target size equals input size in my DQN code?

everyone!When I was doing dqn programming, I encountered some problems. This error says
“ Userwarning: Using a target size (torch.Size([32,32])) that is different to the input size (torch.Size([32,1])).This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.

return F.mse_loss(input,target,reduction=self.reduction)"

And I don't know where the mistake is because I am new to RL . And some of these codes are borrowed from other people's codes, so I don't understand some places.

here are codes:

# hyperparameters
gamma = 0.9
TARGET_REPLACE_ITER = 20
memory_capability = 100    
batch_size = 32
learning_rate = 0.001
n_state = 5
n_action = 32 

neural network code:

class NN(nn.Module):

def __init__(self, ):
    super(NN,self).__init__()
    self.fc1 = nn.Linear(n_state, 32)
    self.fc1.weight.data.normal_(0, 0.1)
    self.fc2 = nn.Linear(32,64)
    self.out = nn.Linear(64, n_action)
    self.out.weight.data.normal_(0, 0.1)

def forward(self, x):
    x = self.fc1(x)
    x = F.relu(x)
    x = self.fc2(x)
    x = F.relu(x)
    action_value = self.out(x)
    return action_value

agent code:

class Agent(object):
    def __init__(self,):
        self.learn_step_counter = 0
        self.memory = np.zeros((memory_capability, n_state * 2 + 2))
        self.memory_cntr = 0
        self.eval_net, self.target_net = NN(), NN()
        self.loss_func = nn.MSELoss()
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=learning_rate)


    def choose_action(self, state):       
        state = torch.unsqueeze(torch.FloatTensor(state),0)   # state is 1-Dim np.array,shape = (5,)
        if random.random() < epsilon:         
            action = random.randint(0,len(stringlist) - 1)
        
        else:
            action_value = self.eval_net.forward(state)
            action = torch.max(action_value, 1)[1].numpy()[0]      
        return action

    def learn(self):   
        if self.learn_step_counter % TARGET_REPLACE_ITER == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        self.learn_step_counter += 1
        sample_index = np.random.choice(memory_capability, batch_size)
        b_memory = self.memory[sample_index, :]
        b_s = torch.FloatTensor(b_memory[:, :n_state])
        b_a = torch.LongTensor(b_memory[:, n_state:n_state + 1].astype(int))
        b_r = torch.FloatTensor(b_memory[:, n_state + 1:n_state + 2])
        b_s_ = torch.FloatTensor(b_memory[:, -n_state:])

        q_eval = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)
        q_next = self.target_net(b_s_).detach()  
        q_target = b_r + gamma * q_next.max(1)[0]  # other people's code said the shape is (batch, 1)=(32,1),but when i ran ,it was (batch,batch)=(32,32),i don't know why
        loss = self.loss_func(q_eval, q_target)        
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def store_transition(self,state,action,reward,state_):
        transition = np.hstack((state,action,reward,state_))
        index = self.memory_cntr % memory_capability
        self.memory[index,:] = transition
        self.memory_cntr += 1

the problem is probably in learn(),but i don't know how to modify.I will appreciate it if someone can help me,thanks a lot

The bug is exactly at the line you pointed out: q_target = b_r + gamma * q_next.max(1)[0]

Here q_next is of shape [batch_size, n_action], so q_next.max(1)[0] is of shape [batch_size]. We also have b_r with a shape of [batch_size,1]. Now adding those two entities will not throw an error as PyTorch is doing some automatic shape broadcasting. So the fix for this is to reshape b_r to [batch_size] from [batch_size,1] by using b_r.squeeze(1)