https://github.com/marrlab/DomainLab/blob/3eed76644bd69b01f6788c7b097a6015285e0394/domainlab/algos/trainers/a_trainer.py#L322