Solved by shengchunnan in post #3 
              
              
                I updated the code a little and now it works. 
Many thanks for your kind suggestion. 
import torch
import torch.nn as nn
import torch.fx as fx

class Network(nn.Module):
    def __init__(self, num_classes=10):
        super(Network, self).__init__()

        self.conv1 = nn.Conv2d(1, 8, kernel_size=…