trainAE.py

    def trainModel(self):
        for epoch in range(self.epochs):
            for data in self.dataLoader:
                image, _ = data
                image = image.view(image.size(0), -1)
                image = Variable(image)
                # Predictions
                output = self(image)
                # Calculate Loss
                loss = self.criterion(output, image)
                # Backpropagation
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            print('epoch [{}/{}], loss:{:.4f}'
                  .format(epoch + 1, self.epochs, loss.data))