From 292af1f2f4eaa4158c612f6e09a0d9d49d04e8dc Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 24 Sep 2018 03:10:42 +0200
Subject: [PATCH] align loss to darknet

---
 models.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/models.py b/models.py
index 8414d4e4..59207bb4 100755
--- a/models.py
+++ b/models.py
@@ -137,8 +137,8 @@ class YOLOLayer(nn.Module):
 
         # Training
         if targets is not None:
-            MSELoss = nn.MSELoss()
-            BCEWithLogitsLoss = nn.BCEWithLogitsLoss()
+            MSELoss = nn.MSELoss(size_average=False)
+            BCEWithLogitsLoss = nn.BCEWithLogitsLoss(size_average=False)
             # CrossEntropyLoss = nn.CrossEntropyLoss()
 
             if requestPrecision:
@@ -160,7 +160,7 @@ class YOLOLayer(nn.Module):
             nT = sum([len(x) for x in targets])  # number of targets
             nM = mask.sum().float()  # number of anchors (assigned to targets)
             nB = len(targets)  # batch size
-            k = nM / nB
+            k = 1 / nB
             if nM > 0:
                 lx = k * MSELoss(x[mask], tx[mask])
                 ly = k * MSELoss(y[mask], ty[mask])