From 5170cd36b0a6c51e8f93f0aa2fda15f874d625dc Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 3 Apr 2019 11:31:31 +0200
Subject: [PATCH] updates

---
 train.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/train.py b/train.py
index f5712b21a3..5979626c3f 100644
--- a/train.py
+++ b/train.py
@@ -40,8 +40,6 @@ def train(
 
     # Initialize model
     model = Darknet(cfg, img_size).to(device)
-    # for m in model.modules():
-    #     weights_init_normal(m)  # set weight distributions
 
     # Optimizer
     lr0 = 0.001  # initial learning rate
@@ -56,7 +54,7 @@ def train(
     if resume:  # Load previously saved model
         if transfer:  # Transfer learning
             chkpt = torch.load(weights + 'yolov3.pt', map_location=device)
-            model.load_state_dict({k: v for k, v in chkpt['model'].items() if v.numel() > 1 and v.shape[0] != nf},
+            model.load_state_dict({k: v for k, v in chkpt['model'].items() if v.numel() > 1 and v.shape[0] != 255},
                                   strict=False)
             for p in model.parameters():
                 p.requires_grad = True if p.shape[0] == nf else False