diff --git a/main.py b/main.py
index beb9f4d..5b6cd80 100644
--- a/main.py
+++ b/main.py
@@ -14,7 +14,7 @@ import numpy as np
 import time
 from tensorboardX import SummaryWriter
 from datasets import __datasets__
-from models import __models__
+from models import __models__, model_loss
 from utils import *
 from torch.utils.data import DataLoader
 import gc