Fix multiprocessing pickling issue