diff --git a/api/src/main.py b/api/src/main.py
index f3787f0..ea02671 100644
--- a/api/src/main.py
+++ b/api/src/main.py
@@ -1,3 +1,4 @@
+from r_cnn.r_cnn_test import train_cnn_test
 from yolo.train_yolo_faces import train_yolo
 
 
diff --git a/api/src/r_cnn/r_cnn.py b/api/src/r_cnn/r_cnn.py
index 73a4c03..199b135 100644
--- a/api/src/r_cnn/r_cnn.py
+++ b/api/src/r_cnn/r_cnn.py
@@ -4,7 +4,7 @@ import torch
 import torchvision.transforms.functional as TF
 from tqdm import tqdm
 import torch.nn as nn
-from src.util import iou, visualizeImage
+from util import iou, visualizeImage
 import random
 import torch.nn.functional as F
 
diff --git a/api/src/r_cnn/r_cnn_test.py b/api/src/r_cnn/r_cnn_test.py
index 6d5e86b..de3ab82 100644
--- a/api/src/r_cnn/r_cnn_test.py
+++ b/api/src/r_cnn/r_cnn_test.py
@@ -8,40 +8,43 @@ import torch.nn as nn
 from .r_cnn import ObjectDetectionCNN, train, eval
 from .cocoDetectionDataset import CocoDetectionDataset
 
-SAVE_PATH = "./saved_models"
-PADDING = 20
 
 def get_transform():
     return ToTensor()
 
-train_dataset = CocoDetectionDataset(
-    image_dir="data/football/train", 
-    annotation_path="data/football/train/_annotations.coco.json",
-    transforms=get_transform()
-)
+def train_cnn_test():
 
-val_dataset = CocoDetectionDataset(
-    image_dir="data/football/valid",
-    annotation_path="data/football/valid/_annotations.coco.json",
-    transforms=get_transform()
-)
+    SAVE_PATH = "./saved_models"
+    PADDING = 20
 
-train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
-val_loader = DataLoader(val_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
+    train_dataset = CocoDetectionDataset(
+        image_dir="data/football/train", 
+        annotation_path="data/football/train/_annotations.coco.json",
+        transforms=get_transform()
+    )
 
-device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
-print("Using device", device)
+    val_dataset = CocoDetectionDataset(
+        image_dir="data/football/valid",
+        annotation_path="data/football/valid/_annotations.coco.json",
+        transforms=get_transform()
+    )
 
-model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
-model.to(device)
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
-loss_module = nn.CrossEntropyLoss()
+    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
+    val_loader = DataLoader(val_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
 
-#train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
-#                optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, PADDING=40, device=device)
+    device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
+    print("Using device", device)
 
-#exit()
+    model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
+    model.to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
+    loss_module = nn.CrossEntropyLoss()
 
-image = next(iter(val_loader))[0][0]
-eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "object_detection", "object_detection"), 
-     device=device, PADDING=40, minSize=5, maxSize=100, minConf=0.8)
\ No newline at end of file
+    #train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
+    #                optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, PADDING=40, device=device)
+
+    #exit()
+
+    image = next(iter(val_loader))[0][0]
+    eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "object_detection", "object_detection"), 
+        device=device, PADDING=40, minSize=5, maxSize=100, minConf=0.8)
\ No newline at end of file
diff --git a/api/src/yolo/yolo_loss.py b/api/src/yolo/yolo_loss.py
index a53f108..d27a572 100644
--- a/api/src/yolo/yolo_loss.py
+++ b/api/src/yolo/yolo_loss.py
@@ -12,13 +12,16 @@ class YoloLoss(nn.Module):
         target_boxes = targets[..., :4]
         target_conf = targets[..., 4]
         target_classes = targets[..., 5:]
+
+        obj_mask = targets[..., 4] == 1
+        noobj_mask = targets[..., 4] == 0
         
-        box_loss = lambda_coord * torch.mean((pred_boxes - target_boxes) ** 2)
+        box_loss = lambda_coord * torch.mean((pred_boxes[obj_mask] - target_boxes[obj_mask]) ** 2)
 
-        obj_loss = torch.mean((pred_conf[target_conf == 1] - target_conf[target_conf == 1]) ** 2)
-        noobj_loss = lambda_noobj * torch.mean((pred_conf[target_conf == 0]) ** 2)
+        obj_loss = torch.mean((pred_conf[obj_mask] - target_conf[obj_mask]) ** 2)
+        noobj_loss = lambda_noobj * torch.mean((pred_conf[noobj_mask]) ** 2)
 
-        class_loss = torch.mean((pred_classes[target_conf == 1] - target_classes[target_conf == 1]) ** 2)
+        class_loss = torch.mean((pred_classes[obj_mask] - target_classes[obj_mask]) ** 2)
 
         total_loss = box_loss + obj_loss + noobj_loss + class_loss
         return total_loss
diff --git a/api/src/yolo/yolo_model.py b/api/src/yolo/yolo_model.py
index cdbbf0c..76bf15b 100644
--- a/api/src/yolo/yolo_model.py
+++ b/api/src/yolo/yolo_model.py
@@ -30,7 +30,8 @@ class Yolo_model(nn.Module):
             nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(),
             nn.Flatten(),
-            nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels))
+            nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels)),
+            nn.Dropout(0.3)
         )
 
     def forward(self, x):