diff --git a/ctapipe/tools/train_disp_reconstructor.py b/ctapipe/tools/train_disp_reconstructor.py index 767e86c2729..5e0c88c9a8a 100644 --- a/ctapipe/tools/train_disp_reconstructor.py +++ b/ctapipe/tools/train_disp_reconstructor.py @@ -73,14 +73,16 @@ def setup(self): """ Initialize components from config. """ - self.loader = TableLoader( - parent=self, - load_dl1_images=False, - load_dl1_parameters=True, - load_dl2=True, - load_simulated=True, - load_instrument=True, - load_observation_info=True, + self.loader = self.enter_context( + TableLoader( + parent=self, + load_dl1_images=False, + load_dl1_parameters=True, + load_dl2=True, + load_simulated=True, + load_instrument=True, + load_observation_info=True, + ) ) self.n_events.attach_subarray(self.loader.subarray) diff --git a/ctapipe/tools/train_energy_regressor.py b/ctapipe/tools/train_energy_regressor.py index dd0f16b5121..b022507e668 100644 --- a/ctapipe/tools/train_energy_regressor.py +++ b/ctapipe/tools/train_energy_regressor.py @@ -73,13 +73,15 @@ def setup(self): """ Initialize components from config """ - self.loader = TableLoader( - parent=self, - load_dl1_images=False, - load_dl1_parameters=True, - load_dl2=True, - load_simulated=True, - load_instrument=True, + self.loader = self.enter_context( + TableLoader( + parent=self, + load_dl1_images=False, + load_dl1_parameters=True, + load_dl2=True, + load_simulated=True, + load_instrument=True, + ) ) self.n_events.attach_subarray(self.loader.subarray) diff --git a/ctapipe/tools/train_particle_classifier.py b/ctapipe/tools/train_particle_classifier.py index d3a1b9d90b5..5c7a8852d50 100644 --- a/ctapipe/tools/train_particle_classifier.py +++ b/ctapipe/tools/train_particle_classifier.py @@ -102,24 +102,28 @@ def setup(self): Initialize components from config """ - self.signal_loader = TableLoader( - parent=self, - input_url=self.input_url_signal, - load_dl1_images=False, - load_dl1_parameters=True, - load_dl2=True, - load_simulated=True, - load_instrument=True, + self.signal_loader = self.enter_context( + TableLoader( + parent=self, + input_url=self.input_url_signal, + load_dl1_images=False, + load_dl1_parameters=True, + load_dl2=True, + load_simulated=True, + load_instrument=True, + ) ) - self.background_loader = TableLoader( - parent=self, - input_url=self.input_url_background, - load_dl1_images=False, - load_dl1_parameters=True, - load_dl2=True, - load_simulated=True, - load_instrument=True, + self.background_loader = self.enter_context( + TableLoader( + parent=self, + input_url=self.input_url_background, + load_dl1_images=False, + load_dl1_parameters=True, + load_dl2=True, + load_simulated=True, + load_instrument=True, + ) ) if self.signal_loader.subarray != self.background_loader.subarray: