Function AddmmBackward returned an invalid gradient
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 3)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = NeuralNetwork()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split

def UploadData(path, train):
    #set up transforms for train and test datasets
    train_transforms = transforms.Compose([transforms.Grayscale(num_output_channels=1), transforms.Resize(255), transforms.CenterCrop(224), transforms.RandomRotation(30), 
                                         transforms.RandomHorizontalFlip(), transforms.transforms.ToTensor()]) 
    valid_transforms = transforms.Compose([transforms.Grayscale(num_output_channels=1), transforms.Resize(255), transforms.CenterCrop(224), transforms.RandomRotation(30), 
                                         transforms.RandomHorizontalFlip(), transforms.transforms.ToTensor()]) 
    test_transforms = transforms.Compose([transforms.Grayscale(num_output_channels=1), transforms.Resize(255), transforms.CenterCrop(224), transforms.ToTensor()])
    
    #set up datasets from Image Folders
    train_dataset = datasets.ImageFolder(path + '/train', transform=train_transforms)
    valid_dataset = datasets.ImageFolder(path + '/validation', transform=valid_transforms)
    test_dataset = datasets.ImageFolder(path + '/test', transform=test_transforms)

    #set up dataloaders with batch size of 32
    trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
    validloader = torch.utils.data.DataLoader(valid_dataset, batch_size=32, shuffle=True)
    testloader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=True)
  
    return trainloader, validloader, testloader

trainloader, validloader, testloader = UploadData(&quot;/home/lns/research/dataset&quot;, True)

epochs = 5
min_valid_loss = np.inf
for e in range(epochs):
    train_loss = 0.0
    for data, labels in trainloader:
        # Transfer Data to GPU if available
        if torch.cuda.is_available():
            print(&quot;using GPU for data&quot;)
            data, labels = data.cuda(), labels.cuda()
         
        # Clear the gradients
        optimizer.zero_grad()
        # Forward Pass
        target = net(data)
        # Find the Loss
        loss = criterion(target,labels)
        # Calculate gradients
        loss.backward()
        # Update Weights
        optimizer.step()
        # Calculate Loss
        train_loss += loss.item()
     
    valid_loss = 0.0
    model.eval()     # Optional when not using Model Specific layer
    for data, labels in validloader:
        # Transfer Data to GPU if available
        if torch.cuda.is_available():
            print(&quot;using GPU for data&quot;)
            data, labels = data.cuda(), labels.cuda()
         
        # Forward Pass
        target = net(data)
        # Find the Loss
        loss = criterion(target,labels)
        # Calculate Loss
        valid_loss += loss.item()
 
    print('Epoch ',e+1, '\t\t Training Loss: ',train_loss / len(trainloader),' \t\t Validation Loss: ',valid_loss / len(validloader))
     
    if min_valid_loss &gt; valid_loss:
        print(&quot;Validation Loss Decreased(&quot;,min_valid_loss,&quot;---&gt;&quot;,valid_loss,&quot;) \t Saving The Model&quot;)
        min_valid_loss = valid_loss
         
        # Saving State Dict
        torch.save(net.state_dict(), '/home/lns/research/MODEL.pth')



After searching a lot i am asking for help. Can someone help me
understand why this error is occuring in backward propagation.
i followed pytorch cnn tutorail and geeksforgeeks tutorial
dataset is x ray images transformed into grayscale and resize to 255
Is my neural network is wrong or data is not processed correctly?


This is a size mismmatch between the output of your CNN and the number of neurons on on your first fully-connected layer. Because of missing padding, the number of elements when flattened is 16*4*4 i.e. 256 (and not 16*5*5):
self.fc1 = nn.Linear(256, 120)

Once modified, the model will run correctly:
&gt;&gt;&gt; model = NeuralNetwork()
&gt;&gt;&gt; model(torch.rand(1, 1, 28, 28)).shape
torch.Size([1, 3])


Alternatively, you can use an nn.LazyLinear which will deduce the in_feature argument during the very first inference based on its input shape.
self.fc1 = nn.LazyLinear(120)

