From cf2a0cc6d6ea1d35d4647ba43ce10b7216b92fcf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=A9tan=20Bahl?= Date: Tue, 3 May 2022 01:00:38 +0200 Subject: [PATCH] Fix error when using other than 4 channel images with BYOL task (#522) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: BAHL Gaetan --- torchgeo/trainers/byol.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchgeo/trainers/byol.py b/torchgeo/trainers/byol.py index 5080c7128cf..1a2fca522ef 100644 --- a/torchgeo/trainers/byol.py +++ b/torchgeo/trainers/byol.py @@ -347,7 +347,7 @@ def config_task(self) -> None: ) encoder.conv1 = new_layer - self.model = BYOL(encoder, image_size=(256, 256)) + self.model = BYOL(encoder, in_channels=in_channels, image_size=(256, 256)) def __init__(self, **kwargs: Any) -> None: """Initialize a LightningModule for pre-training a model with BYOL.