fixes to squeeze excite and stochastic depth

openclimatefix · Oct 29, 2023 · 2dce68d · 2dce68d
1 parent 0ba8289
commit 2dce68d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/metnet/layers/SqueezeExcitation.py b/metnet/layers/SqueezeExcitation.py
@@ -56,7 +56,7 @@ def forward(self, X: torch.Tensor) -> torch.Tensor:
         torch.Tensor
             Output Tensor
         """
-        x_se = X.mean((2, 3), keepdim=True)
+        x_se = X.mean((2, 3), keepdim=True)  # Mean along H, W dim
         x_se = self.conv_reduce(x_se)
         x_se = self.act1(x_se)
         x_se = self.conv_expand(x_se)

diff --git a/metnet/layers/StochasticDepth.py b/metnet/layers/StochasticDepth.py
@@ -22,7 +22,7 @@ def __init__(self, drop_prob: float = 0.0) -> None:
             probability to drop the network path, by default 0.0
         """
         super().__init__()
-        assert 0 < drop_prob < 1.0
+        assert 0 <= drop_prob <= 1.0
         self.drop_prob = drop_prob
 
     def forward(self, X: torch.Tensor) -> torch.Tensor: