Repository: Lightning-Universe/lightning-flash

Issue:
ModuleNotFoundError: No module named 'sahi.model'
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

I just want to run tutorial [here](https://lightning-flash.readthedocs.io/en/stable/reference/instance_segmentation.html).


### Environment
- Ubuntu Desktop 22.04 LTS
- GPU RTX2080
- Driver nvidia-driver-510



### To Reproduce

- Dockerfile
```
FROM nvidia/cuda:11.6.0-cudnn8-devel-ubuntu20.04

# https://serverfault.com/questions/683605/docker-container-time-timezone-will-not-reflect-changes
ENV TZ=Asia/Tokyo
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone

RUN apt-get update && apt-get install -y python3 python3-pip \
&& apt-get install ffmpeg libsm6 libxext6  -y

RUN pip3 install --no-cache-dir \
   jupyterlab \
   lightning-flash[image] \
   icedata \
   icevision
```


- docker-compose.yaml
```
version: "3"
services:
  notebook:
    build: .
    volumes:
      - ".:/home/work"
      - ".jupyter:/root/.jupyter"
    ports:
      - "7777:7777"
    tty: true
    environment:
      - JUPYTER_ENABLE_LAB=yes
      - NVIDIA_VISIBLE_DEVICES=all
      - NVIDIA_DRIVER_CAPABILITIES=all
    deploy:
       resources:
         reservations:
           devices:
             - capabilities:
               - gpu
    command: jupyter lab --ip=0.0.0.0 --port=7777 --allow-root --no-browser --NotebookApp.token=''
```

Then, run
```
docker compose up
```

Connect to `localhost:7777` and open up jupyter.

Then, run
```
from functools import partial

import flash
from flash.core.utilities.imports import example_requires
# from flash.image import InstanceSegmentation, InstanceSegmentationData

example_requires("image")

import icedata  # noqa: E402
```

I got the error

> ModuleNotFoundErrorTraceback (most recent call last)
> Cell In [3], line 9
>       5 # from flash.image import InstanceSegmentation, InstanceSegmentationData
>       7 example_requires("image")
> ----> 9 import icedata  # noqa: E402
> 
> File /usr/local/lib/python3.8/dist-packages/icedata/__init__.py:1
> ----> 1 import icedata.utils
>       2 import icedata.template
>       4 from icedata.utils import *
> 
> File /usr/local/lib/python3.8/dist-packages/icedata/utils/__init__.py:1
> ----> 1 from icedata.utils.utils import *
>       2 from icedata.utils.plot_size_histogram import *
> 
> File /usr/local/lib/python3.8/dist-packages/icedata/utils/utils.py:3
>       1 __all__ = ["load_model_weights_from_url", "load_data"]
> ----> 3 from icevision.all import *
>       6 def load_model_weights_from_url(
>       7     model: nn.Module, url: str, map_location=torch.device("cpu"), **kwargs
>       8 ) -> None:
>       9     state_dict = torch.hub.load_state_dict_from_url(
>      10         url, map_location=map_location, **kwargs
>      11     )
> 
> File /usr/local/lib/python3.8/dist-packages/icevision/__init__.py:7
>       5 from icevision.data import *
>       6 from icevision import backbones
> ----> 7 from icevision import models
>       8 from icevision.metrics import *
>       9 from icevision.visualize import *
> 
> File /usr/local/lib/python3.8/dist-packages/icevision/models/__init__.py:32
>      29     from icevision.models import fastai
>      31 if SoftDependencies.sahi:
> ---> 32     from icevision.models import inference_sahi
> 
> File /usr/local/lib/python3.8/dist-packages/icevision/models/inference_sahi.py:14
>      10 from icevision.visualize.utils import *
>      11 from icevision.tfms.albumentations import albumentations_adapter
> ---> 14 from sahi.model import DetectionModel
>      15 from sahi.prediction import ObjectPrediction
>      16 from sahi.predict import get_sliced_prediction as sahi_get_sliced_prediction
> 
> ModuleNotFoundError: No module named 'sahi.model'




Discussion:

Resolving patch:
diff --git a/flash/core/trainer.py b/flash/core/trainer.py
--- a/flash/core/trainer.py
+++ b/flash/core/trainer.py
@@ -187,9 +187,9 @@ def predict(
             Returns a list of dictionaries, one for each provided dataloader containing their respective predictions.
         """
         # Note: Prediction on TPU device with multi cores is not supported yet
-        if isinstance(self.accelerator, TPUAccelerator) and self.tpu_cores > 1:
+        if isinstance(self.accelerator, TPUAccelerator) and self.num_devices > 1:
             raise NotImplementedError(
-                f"Prediction on TPU device with multi-cores (requested cores: {self.tpu_cores}) is not supported yet."
+                f"Prediction on TPU device with multi-cores (requested cores: {self.num_devices}) is not supported yet."
             )
         model = model or self.lightning_module
         output_transform = getattr(model, "_output_transform", None) or OutputTransform()
diff --git a/flash/core/utilities/imports.py b/flash/core/utilities/imports.py
--- a/flash/core/utilities/imports.py
+++ b/flash/core/utilities/imports.py
@@ -87,6 +87,7 @@ class Image:
     _PL_GREATER_EQUAL_1_4_0 = compare_version("pytorch_lightning", operator.ge, "1.4.0")
     _PL_GREATER_EQUAL_1_5_0 = compare_version("pytorch_lightning", operator.ge, "1.5.0")
     _PL_GREATER_EQUAL_1_6_0 = compare_version("pytorch_lightning", operator.ge, "1.6.0rc0")
+    _PL_GREATER_EQUAL_1_8_0 = compare_version("pytorch_lightning", operator.ge, "1.8.0")
     _PANDAS_GREATER_EQUAL_1_3_0 = compare_version("pandas", operator.ge, "1.3.0")
     _ICEVISION_GREATER_EQUAL_0_11_0 = compare_version("icevision", operator.ge, "0.11.0")
     _TM_GREATER_EQUAL_0_7_0 = compare_version("torchmetrics", operator.ge, "0.7.0")


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -44,6 +44,7 @@
     _GRAPH_TESTING,
     _IMAGE_AVAILABLE,
     _IMAGE_TESTING,
+    _PL_GREATER_EQUAL_1_8_0,
     _TABULAR_TESTING,
     _TEXT_TESTING,
     _TORCH_OPTIMIZER_AVAILABLE,
@@ -471,7 +472,10 @@ def train_dataloader(self):
 
     assert task.get_num_training_steps() == batch_count
     assert isinstance(trainer.optimizers[0], torch.optim.Adadelta)
-    assert isinstance(trainer.lr_schedulers[0]["scheduler"], torch.optim.lr_scheduler.LambdaLR)
+    if _PL_GREATER_EQUAL_1_8_0:
+        assert isinstance(trainer.lr_scheduler_configs[0].scheduler, torch.optim.lr_scheduler.LambdaLR)
+    else:
+        assert isinstance(trainer.lr_schedulers[0]["scheduler"], torch.optim.lr_scheduler.LambdaLR)
 
 
 @pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core.")
diff --git a/tests/examples/test_scripts.py b/tests/examples/test_scripts.py
--- a/tests/examples/test_scripts.py
+++ b/tests/examples/test_scripts.py
@@ -120,6 +120,7 @@
             "template.py",
             marks=[
                 pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core."),
+                pytest.mark.skipif(os.name == "posix", reason="Flaky on Mac OS (CI)"),
                 pytest.mark.skipif(sys.version_info >= (3, 9), reason="Undiagnosed segmentation fault in 3.9"),
             ],
         ),
diff --git a/tests/helpers/task_tester.py b/tests/helpers/task_tester.py
--- a/tests/helpers/task_tester.py
+++ b/tests/helpers/task_tester.py
@@ -102,11 +102,9 @@ def _test_jit_trace(self, tmpdir):
     path = os.path.join(tmpdir, "test.pt")
 
     model = self.instantiated_task
-    trainer = self.instantiated_trainer
     model.eval()
 
-    model.trainer = trainer
-    model = torch.jit.trace(model, self.example_forward_input)
+    model = model.to_torchscript(method="trace", example_inputs=self.example_forward_input)
 
     torch.jit.save(model, path)
     model = torch.jit.load(path)