lightning flash master incompatible with kornia <= 0.5

[t-vi]

🐛 Bug

For me (Python 3.9, PyTorch master, Lightning master and Bolts master, kornia 0.5), the vision finetuning example seems to not work. Upgrading kornia solved it.

To Reproduce

import torchvision
from torch import nn

import flash
from flash import Trainer
from flash.core.classification import Labels
from flash.core.finetuning import FreezeUnfreeze
from flash.data.utils import download_data
from flash.vision import ImageClassificationData, ImageClassifier
import pytorch_lightning as pl


download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", "data/")

# 2. Load the data
datamodule = ImageClassificationData.from_folders(
    train_folder="data/hymenoptera_data/train/",
    val_folder="data/hymenoptera_data/val/",
    test_folder="data/hymenoptera_data/test/",
)


next(iter(datamodule.train_dataloader()))

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-1-55d7bafa41b9> in <module>
     21 
     22 
---> 23 next(iter(datamodule.train_dataloader()))

/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py in __next__(self)
    519             if self._sampler_iter is None:
    520                 self._reset()
--> 521             data = self._next_data()
    522             self._num_yielded += 1
    523             if self._dataset_kind == _DatasetKind.Iterable and \

/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py in _next_data(self)
   1201             else:
   1202                 del self._task_info[idx]
-> 1203                 return self._process_data(data)
   1204 
   1205     def _try_put_index(self):

/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py in _process_data(self, data)
   1227         self._try_put_index()
   1228         if isinstance(data, ExceptionWrapper):
-> 1229             data.reraise()
   1230         return data
   1231 

/usr/local/lib/python3.9/dist-packages/torch/_utils.py in reraise(self)
    432             # have message field
    433             raise self.exc_type(message=msg)
--> 434         raise self.exc_type(msg)
    435 
    436 

ValueError: Caught ValueError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
    return self.collate_fn(data)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/flash/data/batch.py", line 161, in forward
    sample = self.per_sample_transform(sample)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/flash/data/batch.py", line 80, in forward
    sample = self.post_tensor_transform(sample)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/flash/data/utils.py", line 185, in forward
    return self.func(*args, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/flash/data/process.py", line 359, in post_tensor_transform
    return self.current_transform(sample)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/flash/data/utils.py", line 185, in forward
    return self.func(*args, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/flash/data/transforms.py", line 45, in forward
    outputs = super().forward(inputs)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/kornia/geometry/transform/affwarp.py", line 619, in forward
    return resize(input, self.size, self.interpolation, align_corners=self.align_corners, side=self.side)
  File "/usr/local/lib/python3.9/dist-packages/kornia/geometry/transform/affwarp.py", line 546, in resize
    return torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
  File "/usr/local/lib/python3.9/dist-packages/torch/nn/functional.py", line 3611, in interpolate
    raise ValueError(
ValueError: size shape must match input shape. Input is 1D, size is 2

Expected behavior

Get a batch of something.

[edgarriba]

@t-vi what flash version are you using, and kornia ?

[t-vi]

Flash master, Kornia 0.5

[edgarriba]

try to upgrade to kornia > 0.5

[t-vi]

Indeed upgrading to master makes it work again. Thank you so much!
Should we close this or retitle to check for newer kornia?