style: format code with Black

ebtorch/nn/architectures.py

@@ -155,109 +155,107 @@
 # Fully-Connected Block, New version
 # Joint work with Davide Roznowicz (https://github.com/DavideRoznowicz)
 class FCBlock(nn.Module):
     def __init__(
         self,
         in_sizes: Union[List[int], tuple],
         out_size: int,
         bias: Optional[Union[List[bool], tuple, bool]] = None,
         activation_fx: Optional[Union[List, nn.ModuleList, nn.Module]] = None,
         dropout: Optional[Union[List[Union[float, bool]], float, bool, tuple]] = None,
         batchnorm: Optional[Union[List[bool], bool, tuple]] = None,
     ) -> None:
         super().__init__()
 
         self.activation_fx = nn.ModuleList()
 
         error_uneven_size: str = (
             "The length of lists of arguments must be the same across them."
         )
-        error_illegal_dropout: str = (
-            "The 'dropout' argument must be either False, a float, or an iterable of floats and/or False."
-        )
+        error_illegal_dropout: str = "The 'dropout' argument must be either False, a float, or an iterable of floats and/or False."
 
         # Default cases
         if bias is None:
             bias = [True] * len(in_sizes)
         if dropout is None:
             dropout = [False] * len(in_sizes)
         if batchnorm is None:
             batchnorm = [True] * (len(in_sizes) - 1) + [False]
         if activation_fx is None:
             for _ in range(len(in_sizes) - 1):
                 self.activation_fx.append(nn.ReLU())
             self.activation_fx.append(nn.Identity())
 
         # Ergonomics
         if isinstance(bias, bool):
             bias = [bias] * len(in_sizes)
         if isinstance(dropout, bool):
             if not dropout:
                 dropout = [False] * len(in_sizes)
             else:
                 raise ValueError(error_illegal_dropout)
         elif isinstance(dropout, float) or (
             isinstance(dropout, int) and (dropout in (0, 1))
         ):
             dropout = [dropout] * len(in_sizes)
         elif not isinstance(dropout, list):
             raise ValueError(error_illegal_dropout)
 
         if isinstance(batchnorm, bool):
             batchnorm = [batchnorm] * len(in_sizes)
 
         if isinstance(activation_fx, list):
             self.activation_fx = nn.ModuleList(copy.deepcopy(activation_fx))
         elif isinstance(activation_fx, nn.Module) and not isinstance(
             activation_fx, nn.ModuleList
         ):
             for _ in enumerate(in_sizes):
                 self.activation_fx.append(copy.deepcopy(activation_fx))
         elif isinstance(activation_fx, nn.ModuleList):
             self.activation_fx = copy.deepcopy(activation_fx)
 
         # Sanitize
         if (
             not len(in_sizes)
             == len(bias)
             == len(self.activation_fx)
             == len(dropout)
             == len(batchnorm)
         ):
             raise ValueError(error_uneven_size)
 
         # Start with an empty module list
         self.module_battery = nn.ModuleList(modules=None)
 
         for layer_idx in range(len(in_sizes) - 1):
             self.module_battery.append(
                 nn.Linear(
                     in_features=in_sizes[layer_idx],
                     out_features=in_sizes[layer_idx + 1],
                     bias=bias[layer_idx],
                 )
             )
             self.module_battery.append(copy.deepcopy(self.activation_fx[layer_idx]))
 
             if batchnorm[layer_idx]:
                 self.module_battery.append(
                     nn.BatchNorm1d(num_features=in_sizes[layer_idx + 1])
                 )
 
             if isinstance(dropout[layer_idx], bool) and dropout[layer_idx]:
                 raise ValueError(error_illegal_dropout)
             if not isinstance(dropout[layer_idx], bool):
                 self.module_battery.append(nn.Dropout(p=dropout[layer_idx]))
 
         self.module_battery.append(
             nn.Linear(in_features=in_sizes[-1], out_features=out_size, bias=bias[-1])
         )
         self.module_battery.append(copy.deepcopy(self.activation_fx[-1]))
         if batchnorm[-1]:
             self.module_battery.append(nn.BatchNorm1d(num_features=out_size))
         if isinstance(dropout[-1], bool) and dropout[-1]:
             raise ValueError(error_illegal_dropout)
         if not isinstance(dropout[-1], bool):
             self.module_battery.append(nn.Dropout(p=dropout[-1]))
 
     def reset_parameters(self) -> None:
         for module in self.modules():
@@ -813,12 +811,14 @@
         self.extract_z: bool = extract_z
         self.extract_mv: bool = extract_mv
 
-    def forward(self, x: Tensor) -> Union[
+    def forward(
+        self, x: Tensor
+    ) -> Union[
         Tensor,
         Tuple[Tensor, Tensor],
         Tuple[Tensor, Tensor, Tensor],
         Tuple[Tensor, Tensor, Tensor, Tensor],
     ]:
         shared: Tensor = self.encoder(x)
         mean: Tensor = self.mean_neck(shared)
         logvar: Tensor = self.logvar_neck(shared)

ebtorch/optim/lookahead.py

@@ -119,10 +119,10 @@ def step(self, closure=None):
                     param_state["cached_params"].copy_(p.data)
                     if self.pullback_momentum == "pullback":
                         internal_momentum = self.optimizer.state[p]["momentum_buffer"]
-                        self.optimizer.state[p]["momentum_buffer"] = (
-                            internal_momentum.mul_(self.la_alpha).add_(
-                                1.0 - self.la_alpha, param_state["cached_mom"]
-                            )
+                        self.optimizer.state[p][
+                            "momentum_buffer"
+                        ] = internal_momentum.mul_(self.la_alpha).add_(
+                            1.0 - self.la_alpha, param_state["cached_mom"]
                         )
                         param_state["cached_mom"] = self.optimizer.state[p][
                             "momentum_buffer"

ebtorch/optim/schedcos.py

@@ -254,7 +254,8 @@ def _get_lr(self, t):
                 )
                 t_i = self.cycle_mul**i * self.t_initial
                 t_curr = (
-                    t - (1 - self.cycle_mul**i) / (1 - self.cycle_mul) * self.t_initial
+                    t
+                    - (1 - self.cycle_mul**i) / (1 - self.cycle_mul) * self.t_initial
                 )
             else:
                 i = t // self.t_initial