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.ReLU(inplace=True),
        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.ReLU(inplace=True),
        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),
        nn.ReLU(inplace=True)
        )
    
    self.layer4 = nn.Sequential(
        nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
        nn.ReLU(inplace=True)
        )
    
    self.layer5 = nn.Sequential(
        nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=3, stride=1)
        )
    
    self.layer6 = nn.Sequential(
        nn.Dropout(p=0.5)
        )
    
    #self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
    
    self.layer7 = nn.Sequential(
        nn.Linear(in_features=4096, out_features=4096, bias=True),
        nn.ReLU()
        )
    
    self.layer8 = nn.Sequential(
        nn.Dropout(p=0.5)
        )
    
    self.layer9 = nn.Sequential(
        nn.Linear(in_features=4096, out_features=4096, bias=True),
        nn.ReLU()
        )
    
    self.layer10 = nn.Sequential(
        nn.Linear(in_features=4096, out_features=1000, bias=True),
        nn.Softmax()
        )
    
    
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):
try:
    
    _total_steps = len(_train_loader)
    
    for _epochs in range(_num_epochs):
        
        for i, (_images, _labels) in enumerate(_train_loader):
            
            _images = _images.to(_device)
            _labels = _labels.to(_device)
            
            # forward pass
            _outputs = _model(_images)
            _loss = _criterion(_outputs, _labels)
        
            # backward pass and optimization
            _optimizer.zero_grad()
            _loss.backward()
            _optimizer.step()
            
            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")
    print(error)
    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!