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
  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")
  
  
  @step
  def pytorch_evaluator(model: nn.Module, test_loader: DataLoader) -> float:
      """Calculate the loss for the model for each epoch in a graph."""
      model.eval()
      test_loss = 0
      correct = 0
      with torch.no_grad():
          for data, target in test_loader:
              data, target = data.to(DEVICE), target.to(DEVICE)
              output = model(data)
              test_loss += F.nll_loss(
                  output, target, reduction="sum"
              ).item()  # sum up batch loss
              pred = output.argmax(
                  dim=1, keepdim=True
              )  # get the index of the max log-probability
              correct += pred.eq(target.view_as(pred)).sum().item()
  
      test_loss /= len(test_loader.dataset)
  
      print(
          "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
              test_loss,
              correct,
              len(test_loader.dataset),
              100.0 * correct / len(test_loader.dataset),
          )
      )
      return correct / len(test_loader.dataset)
  