Update app.py

app.py

@@ -7,6 +7,7 @@ import torch.nn as nn
 import timm
 from torch import nn
 from torchvision import transforms
+import torchvision.models as models
 # Create the 'model' directory if it doesn't exist (for saving the model later)
 if not os.path.exists('model'):
     os.makedirs('model')