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
  from torch import nn
  from torch.utils.data import DataLoader
  
  from zenml import step
  
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
  
  
  class NeuralNetwork(nn.Module):
      def __init__(self):
          super(NeuralNetwork, self).__init__()
          self.flatten = nn.Flatten()
          self.linear_relu_stack = nn.Sequential(
              nn.Linear(28 * 28, 512),
              nn.ReLU(),
              nn.Linear(512, 512),
              nn.ReLU(),
              nn.Linear(512, 10),
          )
  
      def forward(self, x):
          x = self.flatten(x)
          logits = self.linear_relu_stack(x)
          return logits
  
  
  @step(enable_cache=True)
  def trainer(train_dataloader: DataLoader) -> nn.Module:
      """Trains on the train dataloader."""
      model = NeuralNetwork().to(DEVICE)
      loss_fn = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
      size = len(train_dataloader.dataset)
      model.train()
      for batch, (X, y) in enumerate(train_dataloader):
          X, y = X.to(DEVICE), y.to(DEVICE)
  
          # Compute prediction error
          pred = model(X)
          loss = loss_fn(pred, y)
  
          # Backpropagation
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          if batch % 100 == 0:
              loss, current = loss.item(), batch * len(X)
              print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
      return model
  