Error using Pytorch nn.Sequential: RuntimeError: mat1 and mat2 shapes cannot be multiplied

After I try to clean up my working code to use nn.Sequential, I start to get the error when doing a forward pass

> RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x49 and 3136x512)

This is likely because during my original forward pass, I did this in the middle of the pass to get the correct dimensions:

x = x.view(x.size()[0], -1)

How can this be done when we're using nn.Sequential to define the layers?

Old Working:

import torch.nn as nn
import torch.nn.functional as F
import torch

class Net(nn.Module):
    def __init__(self, out_dims):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(4, 32, 8, 4)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(32, 64, 4, 2)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv3 = nn.Conv2d(64, 64, 3, 1)
        self.relu3 = nn.ReLU(inplace=True)
        self.fc4 = nn.Linear(3136, 512)
        self.relu4 = nn.ReLU(inplace=True)
        self.fc5 = nn.Linear(512, out_dims)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.conv3(x)
        x = self.relu3(x)  # torch.Size([1, 64, 7, 7])
        x = x.view(x.size()[0], -1)  # torch.Size([1, 3136])
        x = self.fc4(x)
        x = self.relu4(x)
        x = self.fc5(x)
        
        return x
        
net = Net(2)
x = torch.rand(4, 84, 84).unsqueeze(0)
net(x)

New non-working using nn.Sequential:

import torch.nn as nn
import torch.nn.functional as F
import torch

class Net(nn.Module):
    def __init__(self, out_dims):
        super(Net, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(4, 32, 8, 4),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 4, 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, 1),
            nn.ReLU(inplace=True),
            nn.Linear(3136, 512),  # RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x49 and 3136x512)
            nn.ReLU(inplace=True),
            nn.Linear(512, out_dims),
        )
        
    def forward(self, x):        
        return self.layers(x)
        
net = Net(2)
x = torch.rand(4, 84, 84).unsqueeze(0)
net(x)

You can use the function nn.flatten() in your sequential model to do the same thing:

Before:

import torch.nn as nn
import torch.nn.functional as F
import torch

class Net(nn.Module):
    def __init__(self, out_dims):
        super(Net, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(4, 32, 8, 4),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 4, 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, 1),
            nn.ReLU(inplace=True),
            nn.Linear(3136, 512),  # RuntimeError: mat1 and mat2 shapes cannot 
        be multiplied (64x49 and 3136x512)
            nn.ReLU(inplace=True),
            nn.Linear(512, out_dims),
        )
    
    def forward(self, x):        
        return self.layers(x)
    
net = Net(2)
x = torch.rand(4, 84, 84).unsqueeze(0)
net(x)

Output:

mat1 and mat2 shapes cannot be multiplied (448x7 and 3136x512)

After Adding nn.Flatten():

    class Net(nn.Module):
    def __init__(self, out_dims):
        super(Net, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(4, 32, 8, 4),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 4, 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, 1),
            nn.ReLU(inplace=True),
            nn.Flatten(),
            nn.Linear(3136, 512),  
            nn.ReLU(inplace=True),
            nn.Linear(512, out_dims),
        )
    
    def forward(self, x):        
        return self.layers(x)
    
net = Net(2)
x = torch.rand(4, 84, 84).unsqueeze(0)
net(x)

Output:

tensor([[ 0.0314, -0.0011]], grad_fn=<AddmmBackward0>)