Move loss calculation to head (#2874)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
Ayush Chaurasia 2023-05-30 22:54:30 +05:30 committed by GitHub
parent 7f077f7654
commit facb7861cf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 417 additions and 348 deletions

View file

@ -157,12 +157,12 @@ class BaseValidator:
# Inference
with dt[1]:
preds = model(batch['img'])
preds = model(batch['img'], augment=self.args.augment)
# Loss
with dt[2]:
if self.training:
self.loss += trainer.criterion(preds, batch)[1]
self.loss += model.loss(batch, preds)[1]
# Postprocess
with dt[3]: