Fix training cost logging