Bug fix in trainer