diff --git a/ml_pytorch_mlflow.py b/ml_pytorch_mlflow.py
index 474c72e..598c67a 100644
--- a/ml_pytorch_mlflow.py
+++ b/ml_pytorch_mlflow.py
@@ -4,6 +4,7 @@ import torchvision
 import matplotlib
 import torch.nn as nn
 import pandas as pd
+import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
 import torch.nn.functional as F