Skip to content

Commit

Permalink
Fix error when using other than 4 channel images with BYOL task (#522)
Browse files Browse the repository at this point in the history
Co-authored-by: BAHL Gaetan <[email protected]>
  • Loading branch information
gaetanbahl and BAHL Gaetan authored May 2, 2022
1 parent 8b3e581 commit cf2a0cc
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion torchgeo/trainers/byol.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit cf2a0cc

Please sign in to comment.