Hyperparameter Tuning using MLflow

import torch
import matplotlib.pyplot as plt
import numpy as np

import torchvision
from torchvision import datasets, transforms

import torch.nn as nn
import torch.optim as optim
import seaborn as sns

from ipywidgets import IntProgress

import mlflow
import mlflow.pytorch

class Params(object):
    def __init__(self, batch_size, epochs, seed, log_interval):
        self.batch_size = batch_size
        self.epochs = epochs
        self.seed = seed
        self.log_interval = log_interval

args = Params(256, 2, 0, 20)

# all the thing that you want to change in code - hyperparameters here