Update model.py

model.py

@@ -8,6 +8,7 @@ from torchvision.utils import make_grid, save_image
 import torch.nn.functional as F
 
 import pytorch_lightning as pl
+import numpy as np
 
 import mlflow.pytorch