diff --git a/model/loss.py b/model/loss.py
index e7454e9..72e5de6 100644
--- a/model/loss.py
+++ b/model/loss.py
@@ -2,6 +2,7 @@ import torch
 import numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
+import torchvision.models as models
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")