Make activation function configurable as argument

ice_station_zebra/models/common/bottleneckblock.py

@@ -1,3 +1,5 @@
+from typing import Type
+
 import torch.nn as nn
 from torch import Tensor
 
@@ -9,18 +11,19 @@ def __init__(
         out_channels: int,
         *,
         filter_size: int,
+        activation: Type[nn.Module] = nn.ReLU,
     ) -> None:
         super().__init__()
 
         self.model = nn.Sequential(
             nn.Conv2d(
                 in_channels, out_channels, kernel_size=filter_size, padding="same"
             ),
-            nn.ReLU(inplace=True),
+            activation(),
             nn.Conv2d(
                 out_channels, out_channels, kernel_size=filter_size, padding="same"
             ),
-            nn.ReLU(inplace=True),
+            activation(),
             nn.BatchNorm2d(num_features=out_channels),
         )
 

ice_station_zebra/models/common/convblock.py

@@ -1,3 +1,5 @@
+from typing import Type
+
 import torch.nn as nn
 from torch import Tensor
 
@@ -10,18 +12,19 @@ def __init__(
         *,
         filter_size: int,
         final: bool = False,
+        activation: Type[nn.Module] = nn.ReLU,
     ) -> None:
         super().__init__()
 
         layers = [
             nn.Conv2d(
                 in_channels, out_channels, kernel_size=filter_size, padding="same"
             ),
-            nn.ReLU(inplace=True),
+            activation(),
             nn.Conv2d(
                 out_channels, out_channels, kernel_size=filter_size, padding="same"
             ),
-            nn.ReLU(inplace=True),
+            activation(),
         ]
         if final:
             layers += [
@@ -31,7 +34,7 @@ def __init__(
                     kernel_size=filter_size,
                     padding="same",
                 ),
-                nn.ReLU(inplace=True),
+                activation(),
             ]
 
         else:

ice_station_zebra/models/common/upconvblock.py

@@ -1,15 +1,22 @@
+from typing import Type
+
 import torch.nn as nn
 from torch import Tensor
 
 
 class UpconvBlock(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int) -> None:
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        activation: Type[nn.Module] = nn.ReLU,
+    ) -> None:
         super().__init__()
 
         self.model = nn.Sequential(
             nn.Upsample(scale_factor=2, mode="nearest"),
             nn.Conv2d(in_channels, out_channels, kernel_size=2, padding="same"),
-            nn.ReLU(inplace=True),
+            activation(),
         )
 
     def forward(self, x: Tensor) -> Tensor: