fix: lint

Signed-off-by: Avik Basu <[email protected]>
numaproj · Jun 12, 2024 · b06c340 · b06c340
1 parent 4b84bf2
commit b06c340
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/examples/raw_pivoted.csv2.gz b/examples/raw_pivoted.csv2.gz
diff --git a/numalogic/tools/layer.py b/numalogic/tools/layer.py
@@ -65,6 +65,18 @@ def forward(self, input_: Tensor) -> Tensor:
 
 
 class MultiChannelLinear(nn.Module):
+    """
+    Linear layer that treats each feature as independent isolated channels.
+
+    Args:
+    ----
+        in_features: num of input features
+        out_features: num of output features
+        n_channels: num of channels
+        device: device to run on
+        dtype: datatype to use
+    """
+
     def __init__(
         self, in_features: int, out_features: int, n_channels: int, device=None, dtype=None
     ):
@@ -92,4 +104,7 @@ def forward(self, x):
         return torch.swapdims(output, 0, 1)
 
     def extra_repr(self) -> str:
-        return f"in_features={self.in_features}, out_features={self.out_features}, n_channels={self.n_channels}"
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"n_channels={self.n_channels}"
+        )