content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from rich import print
  from torch.utils.data import DataLoader
  
  from zenml import step
  
  DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
  
  class Net(nn.Module):
      """Network architecture."""
  
      def __init__(self):
          super(Net, self).__init__()
          self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
          self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
          self.conv2_drop = nn.Dropout2d()
          self.fc1 = nn.Linear(320, 50)
          self.fc2 = nn.Linear(50, 10)
  
      def forward(self, x):
          x = F.relu(F.max_pool2d(self.conv1(x), 2))
          x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
          x = x.view(-1, 320)
          x = F.relu(self.fc1(x))
          x = F.dropout(x, training=self.training)
          x = self.fc2(x)
          return F.log_softmax(x, dim=1)
  
  
  @step
  def pytorch_trainer(
      train_loader: DataLoader,
      epochs: int = 2,
      lr: float = 0.01,
  ) -> nn.Module:
      """Train a neural net from scratch to recognize MNIST digits return our model or the learner."""
      model = Net().to(DEVICE)
      optimizer = optim.Adadelta(model.parameters(), lr=lr)
  
      for epoch in range(1, epochs + 1):
          model.train()
          for batch_idx, (data, target) in enumerate(train_loader):
              data, target = data.to(DEVICE), target.to(DEVICE)
              optimizer.zero_grad()
              output = model(data)
              loss = F.nll_loss(output, target)
              loss.backward()
              optimizer.step()
          print("Train Epoch: {} \tLoss: {:.6f}".format(epoch, loss.item()))
  
      return model
  