diff --git a/enspara/util/load.py b/enspara/util/load.py index 977e2b6b4..80e19026e 100644 --- a/enspara/util/load.py +++ b/enspara/util/load.py @@ -10,8 +10,6 @@ import numpy as np import mdtraj as md -# from sklearn.externals.joblib import Parallel, delayed - from .. import exception logger = logging.getLogger(__name__) @@ -116,11 +114,13 @@ def load_as_concatenated(filenames, lengths=None, processes=None, logger.debug("Sounding %s trajectories with %s processes.", len(filenames), processes) with mp.Pool(processes=processes) as pool: - lengths = pool.starmap( + proc = pool.starmap_async( sound_trajectory, [(f, kw.get('stride', 1)) for f, kw in zip(filenames, args) if 'frame' not in kw]) + shapes = proc.get() + # trjs with frame are always length 1, add that to lengths now for i, kw in enumerate(args): if 'frame' in kw: