modify vanilla_model functio to support parametrization

deel-ai · Dec 9, 2024 · 16b0059 · 16b0059
1 parent c6bf80f
commit 16b0059
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/deel/torchlip/modules/downsampling.py b/deel/torchlip/modules/downsampling.py
@@ -31,7 +31,7 @@
 
 class InvertibleDownSampling(torch.nn.PixelUnshuffle, LipschitzModule):
     def __init__(self, kernel_size: int, k_coef_lip: float = 1.0):
-        torch.nn.PixelUnshuffle.__init__(self, downscale_factor  = kernel_size)
+        torch.nn.PixelUnshuffle.__init__(self, downscale_factor=kernel_size)
         LipschitzModule.__init__(self, k_coef_lip)
 
     def vanilla_export(self):

diff --git a/deel/torchlip/modules/module.py b/deel/torchlip/modules/module.py
@@ -72,13 +72,12 @@ def vanilla_model(model: nn.Module):
         model (nn.Module): Lipschitz neural network
     """
     for n, module in model.named_children():
-        if len(list(module.children())) > 0:
-            # compound module, go inside it
-            vanilla_model(module)
-
         if isinstance(module, LipschitzModule):
             # simple module
             setattr(model, n, module.vanilla_export())
+        elif len(list(module.children())) > 0:
+            # compound module, go inside it
+            vanilla_model(module)
 
 
 class _LipschitzCoefMultiplication(nn.Module):

diff --git a/deel/torchlip/modules/upsampling.py b/deel/torchlip/modules/upsampling.py
@@ -32,7 +32,7 @@
 
 class InvertibleUpSampling(torch.nn.PixelShuffle, LipschitzModule):
     def __init__(self, kernel_size: int, k_coef_lip: float = 1.0):
-        torch.nn.PixelShuffle.__init__(self, upscale_factor = kernel_size)
+        torch.nn.PixelShuffle.__init__(self, upscale_factor=kernel_size)
         LipschitzModule.__init__(self, k_coef_lip)
 
     def vanilla_export(self):