From 8ef49f256094f5b5c40302c09f89a69338a9c42f Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 23 Aug 2019 15:27:29 +0200
Subject: [PATCH] updates

---
 train.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index ad13f4c0..9175d4cf 100644
--- a/train.py
+++ b/train.py
@@ -104,7 +104,9 @@ def train():
         # load model
         if opt.transfer:
             chkpt['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
-        model.load_state_dict(chkpt['model'], strict=False)
+            model.load_state_dict(chkpt['model'], strict=False)
+        else:
+            model.load_state_dict(chkpt['model'])
 
         # load optimizer
         if chkpt['optimizer'] is not None: