kwargs for the sub-autoencoders of MultiScaleAutoencoder

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>
IBM · Jul 21, 2023 · abd5f12 · abd5f12
1 parent b28f6de
commit abd5f12
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/simulai/models/_pytorch_models/_autoencoder.py b/simulai/models/_pytorch_models/_autoencoder.py
@@ -1779,7 +1779,8 @@ def __init__(self,
                  shallow:bool=True,
                  scale: float = 1e-3,
                  devices:Union[str, List]="cpu",
-                 name:str=None) -> None:
+                 name:str=None,
+                 **kwargs) -> None:
 
         super(MultiScaleAutoencoder, self).__init__(name=name)
 
@@ -1805,6 +1806,7 @@ def __init__(self,
                         devices=devices,
                         shallow=shallow,
                         name=name,
+                        **kwargs,
                     )
 
             self.AutoencodersList.append(autoencoder)