Fix problem with DataHandler (#92)

Co-authored-by: Florian Kiwit <[email protected]>
QUARK-framework · Nov 23, 2023 · 36e0515 · 36e0515
1 parent f921b4b
commit 36e0515
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/src/modules/applications/QML/generative_modeling/data/data_handler/DataHandler.py b/src/modules/applications/QML/generative_modeling/data/data_handler/DataHandler.py
@@ -125,10 +125,13 @@ def postprocess(self, input_data: dict, config: dict, **kwargs):
             if self.generalization_mark is not None:
                 np.save(f"{store_dir_iter}/histogram_generated.npy", evaluation["histogram_generated"])
             else:
-                samples = input_data["best_sample"]
-                n_shots = np.sum(samples)
-                histogram_generated = np.asarray(samples) / n_shots
-                histogram_generated[histogram_generated == 0] = 1e-8
+                if "best_sample" in list(input_data.keys()):
+                    samples = input_data["best_sample"]
+                    n_shots = np.sum(samples)
+                    histogram_generated = np.asarray(samples) / n_shots
+                    histogram_generated[histogram_generated == 0] = 1e-8
+                else:
+                    histogram_generated = input_data["histogram_generated"]
                 np.save(f"{store_dir_iter}/histogram_generated.npy", histogram_generated)
             self.metrics.add_metric_batch({"histogram_generated": os.path.relpath(
                 f"{store_dir_iter}/histogram_generated.npy_{kwargs['rep_count']}.npy", current_directory)})