text: import numpy as np
  import torch.nn as nn
  import torch.optim as optim
  
  class Example(nn.Module):
      def __init__(self):
          super(Example, self).__init__()
          # parameters
          self.n_fft = 256
  
          # build
          self.conv1 = nn.Conv2d(3, 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, 10)
          # train
          self.loss_func = nn.CrossEntropyLoss()
          self.optimizer = optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
  
      def forward(self, ):
          x = self.pool(F.relu(self.conv1(x)))
          x = self.pool(F.relu(self.conv2(x)))
          x = x.view(-1, 16 * 5 * 5)
          x = F.relu(self.fc1(x))
          x = F.relu(self.fc2(x))
          x = self.fc3(x)
          return x
  
      def train(self):
          for epoch in range(2):  # loop over the dataset multiple times
  
              running_loss = 0.0
              for i, data in enumerate(trainloader, 0):
                  # get the inputs; data is a list of [inputs, labels]
                  inputs, labels = data
  
                  # zero the parameter gradients
                  optimizer.zero_grad()
  
                  # forward + backward + optimize
                  output = self(inputs)
                  loss = self.loss_func(output, labels)
                  loss.backward()
                  optimizer.step()
  
                  # print statistics
                  running_loss += loss.item()
                  if i % 2000 == 1999:    # print every 2000 mini-batches
                      print('[%d, %5d] loss: %.3f' %
                            (epoch + 1, i + 1, running_loss / 2000))
                      running_loss = 0.0
  
              print('Finished Training')
  