This commit is contained in:
Glenn Jocher 2019-04-02 14:29:35 +02:00
parent 748ff9b5b9
commit 47400aa066
1 changed files with 2 additions and 1 deletions

View File

@ -53,7 +53,7 @@ def train(
if checkpoint['optimizer'] is not None:
optimizer.load_state_dict(checkpoint['optimizer'])
best_loss = checkpoint['best_loss']
del checkpoint # current, saved
del checkpoint
else: # Initialize model with backbone (optional)
if cfg.endswith('yolov3.cfg'):
@ -180,6 +180,7 @@ def train(
model) is nn.parallel.DistributedDataParallel else model.state_dict(),
'optimizer': optimizer.state_dict()}
torch.save(checkpoint, latest)
del checkpoint
# Save best checkpoint
if best_loss == mloss['total']: