zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

remove ema.ema hasattr(ema, 'module') check

This commit is contained in:
Glenn Jocher 2020-08-22 15:18:39 -07:00
parent 1309400e6b
commit 83dc540b1d
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@ -321,7 +321,7 @@ def train(hyp, opt, device, tb_writer=None):
results, maps, times = test.test(opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=ema.ema.module if hasattr(ema.ema, 'module') else ema.ema,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir)
@ -353,7 +353,7 @@ def train(hyp, opt, device, tb_writer=None):
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'training_results': f.read(),
'model': ema.ema.module if hasattr(ema, 'module') else ema.ema,
'model': ema.ema,
'optimizer': None if final_epoch else optimizer.state_dict()}
# Save last, best and delete