Black reformatting

Author: BloodAxe

pytorch_toolbelt/inference/ensembling.py

@@ -13,7 +13,11 @@ class ApplySoftmaxTo(nn.Module):
     dim: int
 
     def __init__(
-        self, model: nn.Module, output_key: Union[str, int, Iterable[str]] = "logits", dim: int = 1, temperature: float = 1
+        self,
+        model: nn.Module,
+        output_key: Union[str, int, Iterable[str]] = "logits",
+        dim: int = 1,
+        temperature: float = 1,
     ):
         """
         Apply softmax activation on given output(s) of the model

pytorch_toolbelt/utils/torch_utils.py

@@ -86,7 +86,9 @@ def logit(x: torch.Tensor, eps=1e-5) -> torch.Tensor:
     return torch.log(x / (1.0 - x))
 
 
-def count_parameters(model: nn.Module, keys: Optional[Sequence[str]] = None, human_friendly: bool = False) -> Dict[str, int]:
+def count_parameters(
+    model: nn.Module, keys: Optional[Sequence[str]] = None, human_friendly: bool = False
+) -> Dict[str, int]:
     """
     Count number of total and trainable parameters of a model
     :param model: A model