add vanilla export to InvertibleUpsampling class

deel-ai · Oct 9, 2024 · fbedaac · fbedaac
1 parent 4007e77
commit fbedaac
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/deel/torchlip/modules/upsampling.py b/deel/torchlip/modules/upsampling.py
@@ -43,3 +43,6 @@ def __init__(
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         return F.invertible_upsample(input, self.kernel_size) * self._coefficient_lip
+
+    def vanilla_export(self):
+        return self