whenxuan: black .

Author: wwhenxuan

channel_attention/utils.py

@@ -14,7 +14,7 @@ def create_conv_layer(
 ) -> nn.Module:
     """
     Create a convolutional layer based on the number of dimensions.
-    
+
     :param n_dims: The number of dimensions (1 or 2).
     :param in_channels: Number of input channels.
     :param out_channels: Number of output channels.
@@ -23,7 +23,7 @@ def create_conv_layer(
     :param padding: Padding added to both sides of the input. Default is 0.
     :param bias: If True, adds a learnable bias to the output. Default is True.
     :param groups: Number of blocked connections from input channels to output channels. Default is 1.
-    
+
     :return: A convolutional layer (nn.Conv1d or nn.Conv2d).
     """
     if n_dims == 1:

tests.py

@@ -166,15 +166,19 @@ def test_MultiSEAttention(self) -> None:
 
         # Test MultiSEAttention for time series (1D)
         mse_1d = MultiSEAttention(
-            n_dims=1, n_channels=self.n_channels, n_branches=3,
+            n_dims=1,
+            n_channels=self.n_channels,
+            n_branches=3,
         )
         for x in time_series_inputs:
             output = mse_1d(x)
             self.assertEqual(output.shape, x.shape)
 
         # Test MultiSEAttention for images (2D)
         mse_2d = MultiSEAttention(
-            n_dims=2, n_channels=self.n_channels, n_branches=4,
+            n_dims=2,
+            n_channels=self.n_channels,
+            n_branches=4,
         )
         for x in image_inputs:
             output = mse_2d(x)