From 61e051445c5d38a97f466c35274601b26b91edc4 Mon Sep 17 00:00:00 2001 From: Kristian Fossum Date: Wed, 4 Mar 2026 09:23:36 +0100 Subject: [PATCH] Added option for internal parallelization of runs. E.g. NN evaluations on GPUs --- ensemble/ensemble.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ensemble/ensemble.py b/ensemble/ensemble.py index b8622726..c084d43f 100644 --- a/ensemble/ensemble.py +++ b/ensemble/ensemble.py @@ -265,6 +265,12 @@ def calc_prediction(self, enX=None, save_prediction=None): # Parallelization on HPC using SLURM elif self.sim.input_dict.get('hpc', False): # Run prediction in parallel on hpc en_pred = self.run_on_HPC(enX, batch_size=nparallel) + + # Parallellization internal to the simulator (e.g. batch processing on GPU ) + elif self.sim.input_dict.get('parallel_internal', False): + # make a single matrix for each state + batch_enX = {key: np.array([d[key] for d in enX]) for key in enX[0].keys()} # key: (b, state) + en_pred = self.sim.run_fwd_sim(batch_enX, member_i=None) # Parallelization on local machine using p_map else: