bug fix

wjtscn · Feb 24, 2020 · b4ca1d0 · b4ca1d0
1 parent bc5e666
commit b4ca1d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train_yolov3.py b/train_yolov3.py
@@ -595,14 +595,14 @@ def train(net, train_data, train_dataset, val_data, eval_metric, ctx, save_prefi
         if not FLAGS.nd_only:
             net.hybridize()
         for i, batch in enumerate(train_data):
+            batch_size = batch[0].shape[0]
             if FLAGS.val_only:
                 break
             if (time.time()-st)/60 > FLAGS.max_epoch_time:
                 logger.info('Max epoch time of %d minutes reached after completing %d%% of epoch. '
                             'Moving on to next epoch' % (FLAGS.max_epoch_time, int(100*(i/num_batches))))
                 break
 
-            batch_size = batch[0].shape[0]
             if FLAGS.features_dir is not None:
                 f1 = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
                 f2 = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)