have you benched that parallel is useful here? I guess the psd computation is so fast that it's not useful here. Maybe parallel across epochs but I don't feel it is necessary