Fix error when using other than 4 channel images with BYOL task

[gaetanbahl]

To avoid shape errors, we pass the in_channels argument from BYOLTask to BYOL.

[isaaccorley]

Thanks for the catch! LGTM. Could you paste the error message you received so it can be referenced in the future?

[gaetanbahl]

Thanks for the catch! LGTM. Could you paste the error message you received so it can be referenced in the future?

When I tried to train a ResNet50 with BYOLTask on EuroSAT data (13 bands), I had this error:

Traceback (most recent call last):
  File "train_ssl.py", line 14, in <module>
    task = BYOLTask(in_channels=13,
  File "/opt/conda/lib/python3.8/site-packages/torchgeo/trainers/byol.py", line 370, in __init__
    self.config_task()
  File "/opt/conda/lib/python3.8/site-packages/torchgeo/trainers/byol.py", line 350, in config_task
    self.model = BYOL(encoder, image_size=(256, 256))
  File "/opt/conda/lib/python3.8/site-packages/torchgeo/trainers/byol.py", line 288, in __init__
    self.encoder(torch.zeros(2, self.in_channels, *image_size))
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torchgeo/trainers/byol.py", line 228, in forward
    _ = self.model(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torchvision/models/resnet.py", line 283, in forward
    return self._forward_impl(x)
  File "/opt/conda/lib/python3.8/site-packages/torchvision/models/resnet.py", line 266, in _forward_impl
    x = self.conv1(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 447, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 443, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [64, 13, 7, 7], expected input[2, 4, 256, 256] to have 13 channels, but got 4 channels instead

[adamjstewart]

I wonder if this fixes #337

[isaaccorley]

We were basically hardcoding in_channels=4 since we didn't pass it from BYOLTask to BYOL. I think this likely solves the issue using BYOL with datasets that do not use RGBN NAIP imagery.

[adamjstewart]

Unfortunately this didn't fix #337 but this is obviously a step in the right direction. If anyone figures out how to fix #337 please open a PR!