[fix] forgot to fix ae to resnet naming

uhlmanngroup · Sep 30, 2024 · e8a10ab · e8a10ab
1 parent 119d104
commit e8a10ab
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/bioimage_embed/models/bolts/vqvae.py b/bioimage_embed/models/bolts/vqvae.py
@@ -2,6 +2,7 @@
 from transformers.utils import ModelOutput
 from pythae.models.nn import BaseDecoder, BaseEncoder
 from pythae.models import VAEConfig
+from . import resnets
 
 
 class BaseResNetVQVAEEncoder(BaseEncoder):
@@ -47,7 +48,7 @@ def __init__(
     ):
         super(ResNet50VQVAEEncoder, self).__init__(
             model_config,
-            ae.resnet50_encoder,
+            resnets.resnet50_encoder,
             self.enc_out_dim,
             first_conv,
             maxpool1,
@@ -67,7 +68,7 @@ def __init__(
     ):
         super(ResNet18VQVAEEncoder, self).__init__(
             model_config,
-            ae.resnet18_encoder,
+            resnets.resnet18_encoder,
             self.enc_out_dim,
             first_conv,
             maxpool1,
@@ -115,7 +116,7 @@ def __init__(
         **kwargs,
     ):
         super(ResNet50VQVAEDecoder, self).__init__(
-            model_config, ae.resnet50_decoder, first_conv, maxpool1, **kwargs
+            model_config, resnets.resnet50_decoder, first_conv, maxpool1, **kwargs
         )
 
 
@@ -130,5 +131,5 @@ def __init__(
         **kwargs,
     ):
         super(ResNet18VQVAEDecoder, self).__init__(
-            model_config, ae.resnet18_decoder, first_conv, maxpool1, **kwargs
+            model_config, resnets.resnet18_decoder, first_conv, maxpool1, **kwargs
         )