From effd38dfeadaae6c28ee6e9b74ea6c5442cab8fe Mon Sep 17 00:00:00 2001
From: Maximilian Linhoff <maximilian.linhoff@tu-dortmund.de>
Date: Fri, 10 Mar 2023 18:59:45 +0100
Subject: [PATCH] Add missing enter_context for TableLoader in train tools

---
 ctapipe/tools/train_disp_reconstructor.py  | 18 ++++++-----
 ctapipe/tools/train_energy_regressor.py    | 16 +++++-----
 ctapipe/tools/train_particle_classifier.py | 36 ++++++++++++----------
 3 files changed, 39 insertions(+), 31 deletions(-)

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: