support number of channels!=3 in YOLOConfiguration (#806)

Co-authored-by: Kfir Gedalyahu <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Seppo Enarvi <[email protected]> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lightning-Universe · May 4, 2022 · 55cb279 · 55cb279
1 parent 72614f0
commit 55cb279
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pl_bolts/models/detection/yolo/yolo_config.py b/pl_bolts/models/detection/yolo/yolo_config.py
@@ -40,7 +40,9 @@ def get_network(self) -> nn.ModuleList:
             A :class:`~torch.nn.ModuleList` that defines the YOLO network.
         """
         result = nn.ModuleList()
-        num_inputs = [3]  # Number of channels in the input of every layer up to the current layer
+        # Number of channels in the input of every layer up to the current layer,
+        # use channels from configuration or default of 3
+        num_inputs = [self.global_config.get("channels", 3)]
         for layer_config in self.layer_configs:
             config = {**self.global_config, **layer_config}
             module, num_outputs = _create_layer(config, num_inputs)