From f527d30ccd3ed9f183acf013f6c98fb0263d3f3c Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 2 Apr 2019 18:50:55 +0200
Subject: [PATCH] updates

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

diff --git a/train.py b/train.py
index 79dabb92..39ad63ad 100644
--- a/train.py
+++ b/train.py
@@ -70,10 +70,10 @@ def train(
         del chkpt
 
     else:  # Initialize model with backbone (optional)
-        if cfg.endswith('yolov3.cfg'):
-            cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
-        elif cfg.endswith('yolov3-tiny.cfg'):
+        if '-tiny.cfg' in cfg:
             cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
+        else:
+            cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
 
     # Set scheduler (reduce lr at epoch 250)
     scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[250], gamma=0.1, last_epoch=start_epoch - 1)