Update app.py

app.py

@@ -11,6 +11,7 @@ import numpy as np
 from matplotlib.colors import LinearSegmentedColormap
 import pytorch_lightning as pl
 from tqdm import tqdm
+from torchvision import models
 
 # Function Definitions