stardist_trainer_improved