1 files changed, 6 insertions, 6 deletions
diff --git a/main.py b/main.py
index ca87788d..837dea81 100644
--- a/main.py
+++ b/main.py
@@ -6,11 +6,11 @@ from datetime import datetime
 import tensorflow as tf
 
 import hyperparameters as hp
-from models import YourModel, VGGModel
+from losses import YourModel
 from preprocess import Datasets
 from skimage.transform import resize
-from tensorboard_utils import \
-        ImageLabelingLogger, ConfusionMatrixLogger, CustomModelSaver
+# from tensorboard_utils import \
+#         ImageLabelingLogger, ConfusionMatrixLogger, CustomModelSaver
 
 from skimage.io import imread
 from lime import lime_image
@@ -128,9 +128,9 @@ def train(model, datasets, checkpoint_path, logs_path, init_epoch):
         tf.keras.callbacks.TensorBoard(
             log_dir=logs_path,
             update_freq='batch',
-            profile_batch=0),
-        ImageLabelingLogger(logs_path, datasets),
-        CustomModelSaver(checkpoint_path, ARGS.task, hp.max_num_weights)
+            profile_batch=0)
+        # ImageLabelingLogger(logs_path, datasets),
+        # CustomModelSaver(checkpoint_path, ARGS.task, hp.max_num_weights)
     ]
 
     # Include confusion logger in callbacks if flag set