[Trainer] fix SP loss (#44461)

Commit

1 day ago

[Trainer] fix SP loss (#44461) * fix SP loss * better fallback * Apply suggestion from @kashif