fix possible data leakage

2026-08-20 12:10:19 +08:00 · 2020-07-07 17:45:37 +08:00
parent d39df871a4
commit e2cf6bbcc5
1 changed files with 4 additions and 1 deletions
@@ -136,7 +136,10 @@ class LatentModel(nn.Module):

        if target_y is not None:
            dist_post, log_var_post = self._latent_encoder(target_x, target_y)
-            z = dist_post.loc
+            if self.training:
+                z = dist_post.loc
+            else:
+                z = dist_prior.loc
        else:
            z = dist_prior.loc