diff --git a/confs/config_voc.json b/confs/config_voc.json
index 876c7ac..5f40ee6 100644
--- a/confs/config_voc.json
+++ b/confs/config_voc.json
@@ -9,7 +9,7 @@
     "config_path" : {
         "labels":           "models/voc/labels_voc.txt",
         "anchors":          "models/voc/anchors_voc.txt",
-        "arch_plotname":    ""
+        "arch_plotname":    "voc_arch.png"
     },
     "train": {
         "out_model_name":   "yolo_retrained_voc.h5",
diff --git a/net/netloss.py b/net/netloss.py
index 71bf160..b5b426e 100644
--- a/net/netloss.py
+++ b/net/netloss.py
@@ -55,7 +55,7 @@ def __init__(self):
         self.lambda_class = YoloParams.CLASS_SCALE
 
         self.norm = False
-        self.paper_imp = False
+
 
     def coord_loss(self, y_true, y_pred):
         
@@ -71,7 +71,6 @@ def coord_loss(self, y_true, y_pred):
         if self.norm:
             norm_coord = K.sum(K.cast(indicator_coord > 0.0, np.float32))
 
-
         loss_xy = K.sum(K.square(b_xy - b_xy_pred) * indicator_coord, axis=[1,2,3,4])
         #loss_wh = K.sum(K.square(b_wh - b_wh_pred) * indicator_coord, axis=[1,2,3,4])
         loss_wh = K.sum(K.square(K.sqrt(b_wh) - K.sqrt(b_wh_pred)) * indicator_coord, axis=[1,2,3,4])
@@ -97,32 +96,28 @@ def obj_loss(self, y_true, y_pred):
         if self.norm:
             norm_conf = K.sum(K.cast((indicator_obj + indicator_noobj)  > 0.0), np.float32)
 
-        loss_obj = K.sum(K.square(b_o-b_o_pred) * (indicator_obj + indicator_noobj), axis=[1,2,3])
+        indicator_o = indicator_obj + indicator_noobj
+        loss_obj = K.sum(K.square(b_o-b_o_pred) * indicator_o, axis=[1,2,3])
 
         return loss_obj / (norm_conf + EPSILON) / 2
 
 
     def class_loss(self, y_true, y_pred):
 
-        b_class = K.argmax(y_true[..., 5:], axis=-1)
-        b_class_pred = y_pred[..., 5:]
+        p_c_pred = K.softmax(y_pred[..., 5:])
+        p_c = K.one_hot(K.argmax(y_true[..., 5:], axis=-1), YoloParams.NUM_CLASSES)
+        loss_class_arg = K.sum(K.square(p_c - p_c_pred), axis=-1)
+        
+        #b_class = K.argmax(y_true[..., 5:], axis=-1)
+        #b_class_pred = y_pred[..., 5:]
+        #loss_class_arg = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=b_class, logits=b_class_pred)
 
         indicator_class = y_true[..., 4] * K.gather(
             YoloParams.CLASS_WEIGHTS, b_class) * self.lambda_class
 
         norm_class = 1
         if self.norm:
-            norm_class = K.sum(K.cast(indicator_class > 0.0, np.float32))
-
-        loss_class_arg = tf.nn.sparse_softmax_cross_entropy_with_logits(
-            labels=b_class, logits=b_class_pred)
-
-        '''
-        if self.paper_imp:
-            loss_class_arg = 
-
-        '''
-
+            norm_class = K.sum(K.cast(indicator_class > 0.0, np.float32))        
 
         loss_class = K.sum(loss_class_arg * indicator_class, axis=[1,2,3])
 
diff --git a/net/utils.py b/net/utils.py
index 4ff922a..ab48218 100755
--- a/net/utils.py
+++ b/net/utils.py
@@ -252,7 +252,7 @@ def setup_logging(logging_path='logs'):
     log_path = os.path.join(os.getcwd(),logging_path)
     mkdir_p(log_path)
 
-    check_names = lambda y: y if isinstance(y, int) else -1
+    check_names = lambda y: y if y.isdigit() else -1
     get_ind = lambda x: int(check_names(x.split('_')[1]))
     
     run_counter = max(map(get_ind, os.listdir(log_path)), default=-1) + 1
diff --git a/yolov2.py b/yolov2.py
index 38526f2..35b8219 100644
--- a/yolov2.py
+++ b/yolov2.py
@@ -120,6 +120,8 @@ def inference(self, path):
 
             boxes, scores, _, labels = self.inf_model.predict(image.copy())
 
+            print(f, labels)
+
             image = draw_boxes(image, (boxes, scores, labels))
             out_name =  os.path.join(out_path, os.path.basename(f).split('.')[0] + out_fname_mod)           
             cv2.imwrite(out_name, image)