How to resize a batch of images for use with Pytorch Linear Regression?
I am trying to create a simple linear regression neural net for use with batches of images. The input dimensions are [BatchSize, 3, Width, Height]  with the second dimension representing the RGB channels of the input image.
Here is my (broken) attempt at that regression model:
class LinearNet(torch.nn.Module):
    def __init__(self, Chn, W,H, nHidden):
        &quot;&quot;&quot;
        Input: A [BatchSize x Channels x Width x Height] set of images
        Output: A fitted regression model with weights dimension : [Width x Height]
        &quot;&quot;&quot;
        super(LinearNet, self).__init__()
        self.Chn = Chn
        self.W = W
        self.H = H
        self.hidden = torch.nn.Linear(Chn*W*H,nHidden)   # hidden layer
        self.predict = torch.nn.Linear(nHidden, Chn*W*H)   # output layer

    def forward(self, x):
        torch.reshape(x, (-1,self.Chn*self.W*self.H)) # FAILS here
        # x = x.resize(-1,self.Chn*self.W*self.H)  
        x = F.relu(self.hidden(x))      # activation function for hidden layer
        x = self.predict(x)             # linear output
        x = x.resize(-1,self.Chn, self.W,self.H)
        return x

When sending in a batch of images with dimensions [128 x 3 x 96 x 128]  this fails on the indicated line:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (36864x128 and 36864x256)

How should the matrix dimensions be properly manipulated to use these pytorch functions?
Update Based on a (since deleted) comment I have updated the code to use torch.reshape.

Solution 1 As a possible solution, you can get a batch size from input x with x.shape[0] and use it in reshape later
import torch

batch = torch.zeros([128, 3, 96, 128], dtype=torch.float32)

# -1 will compute last dimension automatically
batch_upd = torch.reshape(batch, (batch.shape[0], -1))

print(batch_upd.shape)

Output for this code is
torch.Size([128, 36864])

Solution 2
As another possible solution you can use flatten
batch_upd = batch.flatten(start_dim=1)

will result in the same output
As to your next problem, consider going through the modified forward code:
def forward(self, x):
    x = x.flatten(1)  # shape: [B, C, W, H] -&gt; [B, C*W*H]
    x = F.relu(self.hidden(x))      # activation function for hidden layer
    x = self.predict(x)             # linear output
    x = x.reshape((-1, self.Chn, self.W, self.H)) # shape: [B, C*W*H] -&gt; [B, C, W, H]
    return x

Here is the successful usage example:
ln = LinearNet(3, 96, 128, 256)
batch = torch.zeros((128, 3, 96, 128))
res = ln(batch)
print(res.shape)  # torch.Size([128, 3, 96, 128])

