Hi all,
I am very new to Pytorch and am trying to set up a simple Convolutional Neural Network. I followed some of the online guides and set up this test network:
import torch
class SimpleCNN(torch.nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        img_size = 50 # image size
        s_im_size = int(img_size/2/2) # image size after 2 times maxPool
        out1 = 2
        out2 = 2
        # Conv Layer 1
        self.layer1 = torch.nn.Sequential(
        torch.nn.Conv2d(1, out1, kernel_size=5, stride=1, padding=2),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2)
        )
        # Conv Layer 2
        self.layer2 = torch.nn.Sequential(
        torch.nn.Conv2d(out1, out2, kernel_size=5, stride=1, padding=2),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2, stride=2)
        )  
        self.fc1 = torch.nn.Linear(out2*s_im_size**2, 3*img_size**2)
        self.fc2 = torch.nn.Linear(3*img_size**2, 3*img_size**2)
    def forward(self, x):
        img_size = 50 # image size
        s_im_size = int(img_size/2/2) # image size after 2 times maxPool
        out2 = 2
        out = self.layer1(x)   # Change from (1,50,50) to (5,25,25)
        out = self.layer2(out)   # Change from (5,25,25) to (7,12,12)
        out = out.view(-1,out2*s_im_size**2) # Change from (7,12,12) to (1,7*12*12)
        out = self.fc1(out)
        out = self.fc2(out)
        out = out.view(img_size**2,3)
        return out

model = SimpleCNN()

Where i should input images of shape (50,50) and get an output of shape (50*50,3). All well and good.
But if I change image size to take slightly larger images ,img_size = 320, and call model=SimpleCNN() i get the error:
RuntimeError: [enforce fail at ..\c10\core\CPUAllocator.cpp:72] data. DefaultCPUAllocator: not enough memory: you tried to allocate 377487360000 bytes. Buy new RAM!

Which to my mind doesn’t quite make sense. I must have made a mistake in my network, can anyone help me figure it out?
cheers