test: hf callback metrics logging fix

[MikhailKardash]

Ticket

Description

TRAIN appears to be a misnomer. Expand a conditional to include both TRAIN and TRAIN_AVG.

Test Plan

Checklist

• [ ] Changes have been manually QA'd
• [ ] New features have been approved by the corresponding PM
• [ ] User-facing API changes have the "User-facing API Change" label
• [ ] Release notes have been added as a separate file under docs/release-notes/ See Release Note for details.
• [ ] Licenses have been included for new code which was copied and/or modified from any external code