Deprecation warnings

model_angelo/models/common_modules.py

@@ -213,3 +213,15 @@ def __init__(self):
     def forward(self, x, y):
         s = torch.sigmoid(self.gate)
         return s * x + (1 - s) * y
+
+
+class Upsample(nn.Module):
+    def __init__(self, scale_factor, mode="nearest"):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+    def forward(self, x):
+        return F.interpolate(
+            x, scale_factor=self.scale_factor, mode=self.mode
+        )

model_angelo/models/resnet.py

@@ -9,6 +9,7 @@
     NormalizeStd,
     RegularBlock,
     ResBlock,
+    Upsample
 )
 
 
@@ -62,7 +63,7 @@ def __init__(
             self.main_layers.append(pooling_class(2, stride=2))
         for i in range(self.num_blocks):
             if i == self.num_blocks - 2 and downsample_x:
-                self.main_layers.append(nn.Upsample(scale_factor=2))
+                self.main_layers.append(Upsample(scale_factor=2))
             self.main_layers.append(
                 ResBlock(
                     in_channels=self.g_width,

model_angelo/utils/torch_utils.py

@@ -571,4 +571,3 @@ def compile_if_possible(module: nn.Module) -> nn.Module:
         module = torch.compile(module)
     return module
 
-