Move loss to task heads (#2825)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
This commit is contained in:
Ayush Chaurasia 2023-05-25 16:07:54 +05:30 committed by GitHub
parent f23a03596d
commit d19c5b6ce8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 335 additions and 346 deletions

View file

@ -162,7 +162,8 @@ class BaseValidator:
# Loss
with dt[2]:
if self.training:
self.loss += trainer.criterion(preds, batch)[1]
loss_items = model.loss(batch, preds)
self.loss += loss_items[1]
# Postprocess
with dt[3]: