Cleaning up VAE while I'm here

uhlmanngroup · Jan 8, 2024 · b24d287 · b24d287
1 parent 245769c
commit b24d287
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/bioimage_embed/models/bolts/vae.py b/bioimage_embed/models/bolts/vae.py
@@ -2,6 +2,9 @@
 from pythae.models.base.base_utils import ModelOutput
 from pythae.models.nn import BaseDecoder, BaseEncoder
 
+
+from pythae import models
+from pythae.models import VQVAEConfig, VAEConfig
 from pl_bolts.models import autoencoders
 from pythae.models import VQVAE, VQVAEConfig, VAE, VAEConfig
 
@@ -34,9 +37,7 @@ def forward(self, x):
         output = ModelOutput()
         x = self.encoder(x)
         # x = self.fc1(x)
-        output["embedding"] = self.embedding(x)
-        output["log_covariance"] = self.log_var(x)
-        return output
+        return ModelOutput(embedding=self.embedding(x), log_covariance=self.log_var(x))
 
 
 class ResNet50VAEDecoder(BaseDecoder):