RuntimeError: Given input size: (256x2x2). Calculated output size: (256x0x0). Output size is too small

This is my code:

import torch
import torch.nn as nn

class AlexNet(nn.Module):

def __init__(self, __output_size):
    super(AlexNet, self).__init__()
    self.layer1 = nn.Sequential(
        nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4),
        nn.MaxPool2d(kernel_size=3, stride=2)
    self.layer2 = nn.Sequential(
        nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2),
        nn.MaxPool2d(kernel_size=3, stride=2)
    self.layer3 = nn.Sequential(
        nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1),
    self.layer4 = nn.Sequential(
        nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
    self.layer5 = nn.Sequential(
        nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1),
        nn.MaxPool2d(kernel_size=3, stride=1)
    self.layer6 = nn.Sequential(
    #self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
    self.layer7 = nn.Sequential(
        nn.Linear(in_features=4096, out_features=4096, bias=True),
    self.layer8 = nn.Sequential(
    self.layer9 = nn.Sequential(
        nn.Linear(in_features=4096, out_features=4096, bias=True),
    self.layer10 = nn.Sequential(
        nn.Linear(in_features=4096, out_features=1000, bias=True),
def forward(self, x):
    _output = self.layer1(x)
    _output = self.layer2(_output)
    _output = self.layer3(_output)
    _output = self.layer4(_output)
    _output = self.layer5(_output)
    _output = self.layer6(_output)
   # _output = self.avgpool(_output)
    _output = torch.flatten(_output, 1)
    _output = self.layer7(_output)
    _output = self.layer8(_output)
    _output = self.layer9(_output)
    _output = self.layer10(_output)
    return _output

This is my training part:

def train(_train_loader, _model, _num_epochs, _device, _criterion, _optimizer):

    _total_steps = len(_train_loader)
    for _epochs in range(_num_epochs):
        for i, (_images, _labels) in enumerate(_train_loader):
            _images =
            _labels =
            # forward pass
            _outputs = _model(_images)
            _loss = _criterion(_outputs, _labels)
            # backward pass and optimization
            if (i+1) % 100 == 0:
                print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(_epochs+1, _num_epochs, i+1, _total_steps, _loss.item()))
except Exception as error:
    print("An error occured while training")
    raise error

Initially I was getting:

RuntimeError: Given groups=1, weight of size [96, 3, 11, 11], expected input[100, 1, 28, 28] to have 3 channels, but got 1 channels instead

Then I changed the number of channels to 1, I got the above mentioned error. Please help me out to resolve this error!

The error is raised, if the spatial size of the input (and thus an intermediate activation) is too small for the model architecture.
It seems you are using an input of 28x28 pixels, which would be too small for an AlexNet-like model, which was initially using inputs of 224x224, so you would have to either resize the inputs to a larger size or modify the model (e.g. by removing layers).

Thanks, I understood!