DQN RuntimeError: mat1 and mat2 shapes cannot be multiplied (18x7 and 126x64)

Question

I am trying to implement the DQN algorithm using pytorch. My environment returns an observation that preprocesses it to a tensor of shape torch.Size([1, 2, 9, 7]). An example of the input:

tensor([[[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.1815, 0.2127, 0.2431, 0.3362, 0.3416],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1969, 0.2341],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.5838, 0.0000, 0.0000, 0.4880, 0.7366, 0.7777, 0.7594]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.1815, 0.2127, 0.2431, 0.3362, 0.3416],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1969, 0.2341],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.5838, 0.0000, 0.0000, 0.4880, 0.7366, 0.7777, 0.7594]]]])

I get the issue with shapes RuntimeError: mat1 and mat2 shapes cannot be multiplied (18x7 and 126x64) after calling the act function of my neural net:

class Net(nn.Module):
    def __init__(self,  env:Env):
        super().__init__()
        self.env = env
        in_features = int(np.prod(env.observation_space.shape))
        self.sequen= nn.Sequential(
            nn.Linear(in_features, 64),
            nn.ReLU(),
            nn.Linear(64, env.action_space.n) 
        )

    def forward(self, x):
        return self.sequen(x)

    def act(self, obs):
        q_values = self(obs).to(device) ### the problem is here !!!!!
        max_q_index = torch.argmax(q_values, dim=1)[0]
        action = max_q_index.detach().item()
        return action

Although I understand that 18*7 = 126 which is the input dimension of the first linear layer, I do not know how to fix the problem. Posts like this and this helped to understand the problem but not how to solve it.

UPDATE
The following code seemed to solve the issue both when passing 1 and several observations.

obs = obs.unsqueeze(0)
obs = obs.view(obs.shape[0], -1)

Answer 1

Your model uses Linear-> ReLU -> Linear architecture, where nn.Linear is just matrix multiplication layer. The first linear matrix is of size 126 x 64 and second linear matrix is of size 64 x env.action_space.n.

Given 2 matrices A and B of size mxn and pxq respectively, for matrix multiplication between A and B, n should be equal to p. The issue is in dimension of input when passed for forward inference to the first layer. For matrix multiplication, your input tensor x of size bx2x9x7 should be changed to bx126 to match first linear matrix dimensions 126 x 64, where b is batch size.

Solution:- use resize (more generic) or flatten to fix the tensor size.

x = x.reshape(x.shape[0], -1)

or

x = torch.flatten(x, start_dim=1)