is target dim required?
Is the definition of target dim in experiment_setup.py
necessary or can lead to confusion and wrong behaviour? Check where it goes to and if it is good to have this flexibility. Otherwise remove it from arguments and code it hard.