title: Lightning Flash error with SemanticSegmentationData (NameError: name 'K' is not defined)
body: <p>I'm trying to perform an image segmentation task with Colab and <a href="https://pypi.org/project/lightning-flash/" rel="nofollow noreferrer">Lightning Flash</a>.</p>
  <p>I'm installing Flash with:</p>
  <pre><code>!pip install lightning-flash
  </code></pre>
  <p>I'm trying to instanciate a Lightning Flash SemanticSegmentationData using from_folders method like this:</p>
  <pre><code>datamodule = SemanticSegmentationData.from_folders(
      train_folder=x_train_dir,
      train_target_folder=y_train_dir,
      val_folder=x_valid_dir,
      val_target_folder=y_valid_dir,
      test_folder=x_test_dir,
      test_target_folder=y_test_dir,
      transform_kwargs=dict(image_size=(256, 256)),
      num_classes=1,
      batch_size=16,
  )
  </code></pre>
  <p>But I'm getting this error:</p>
  <pre><code>/usr/local/lib/python3.7/dist-packages/flash/image/segmentation/input_transform.py in train_per_sample_transform(self)
   49             [DataKeys.INPUT, DataKeys.TARGET],
   50             KorniaParallelTransforms(
  ---&gt; 51                 K.geometry.Resize(self.image_size, interpolation=&quot;nearest&quot;), K.augmentation.RandomHorizontalFlip(p=0.5)
   52             ),
   53         )
  NameError: name 'K' is not defined
  </code></pre>
  <p>How can I solve this problem?</p>