diff --git a/simulai/models/_pytorch_models/_autoencoder.py b/simulai/models/_pytorch_models/_autoencoder.py index 0b855262..a2eba818 100644 --- a/simulai/models/_pytorch_models/_autoencoder.py +++ b/simulai/models/_pytorch_models/_autoencoder.py @@ -1098,6 +1098,7 @@ def __init__( scale: float = 1e-3, devices: Union[str, list] = "cpu", name: str = None, + **kwargs, ) -> None: """ Constructor method. @@ -1172,6 +1173,7 @@ def __init__( shallow=shallow, use_batch_norm=use_batch_norm, name=self.name, + **kwargs ) self.encoder = encoder.to(self.device) diff --git a/simulai/templates/_templates.py b/simulai/templates/_templates.py index a21e217c..3939da05 100644 --- a/simulai/templates/_templates.py +++ b/simulai/templates/_templates.py @@ -673,6 +673,7 @@ def cnn_autoencoder_auto( use_batch_norm: bool = False, shallow: bool = False, name: str = None, + **kwargs, ) -> Tuple[NetworkTemplate, ...]: """ @@ -737,7 +738,7 @@ def cnn_autoencoder_auto( autogen_cnn = NetworkInstanceGen( architecture="cnn", dim=case, use_batch_norm=use_batch_norm, - kernel_size=kernel_size, + kernel_size=kernel_size, **kwargs, ) autogen_dense = NetworkInstanceGen(architecture="dense", shallow=shallow) @@ -798,6 +799,7 @@ def autoencoder_auto( use_batch_norm: bool = False, case: str = None, name: str = None, + **kwargs, ) -> Tuple[Union[NetworkTemplate, None], ...]: """ @@ -864,6 +866,7 @@ def autoencoder_auto( shallow=shallow, use_batch_norm=use_batch_norm, name=name, + **kwargs, ) return encoder, decoder, bottleneck_encoder, bottleneck_decoder diff --git a/tests/network/test_template_gen.py b/tests/network/test_template_gen.py index edf5ff20..e42f86f1 100644 --- a/tests/network/test_template_gen.py +++ b/tests/network/test_template_gen.py @@ -372,6 +372,7 @@ def test_autoencoder_kernel_size_shallow(self) -> None: architecture="cnn", case="2d", shallow=True, + padding_mode='replicate', ) estimated_data = autoencoder.eval(input_data=input_data) @@ -392,6 +393,7 @@ def test_autoencoder_multiscaleautoencoder(self) -> None: case="2d", shallow=True, name="model", + padding_mode='replicate', ) estimated_data = autoencoder.reconstruction_forward(input_data=input_data)