Pass correct training stage in compute_metrics (#534)

* Pass correct training stage in CouplingFlow.compute_metrics

* Pass correct training stage in CIF and PointInferenceNetwork

Author: han-ol

bayesflow/experimental/cif/cif.py

@@ -99,7 +99,7 @@ def _inverse(
     def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "training") -> dict[str, Tensor]:
         base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
 
-        elbo = self.log_prob(x, conditions=conditions)
+        elbo = self.log_prob(x, conditions=conditions, training=stage == "training")
 
         loss = -keras.ops.mean(elbo)
 

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -183,7 +183,7 @@ def compute_metrics(
     ) -> dict[str, Tensor]:
         base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
 
-        z, log_density = self(x, conditions=conditions, inverse=False, density=True)
+        z, log_density = self(x, conditions=conditions, inverse=False, density=True, training=stage == "training")
         loss = weighted_mean(-log_density, sample_weight)
 
         return base_metrics | {"loss": loss}

bayesflow/networks/point_inference_network.py

@@ -145,7 +145,7 @@ def call(
     def compute_metrics(
         self, x: Tensor, conditions: Tensor = None, sample_weight: Tensor = None, stage: str = "training"
     ) -> dict[str, Tensor]:
-        output = self(x, conditions)
+        output = self(x, conditions, training=stage == "training")
 
         metrics = {}
         # calculate negative score as mean over all scores