We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Widgeted progress bar remains visible after suite run
**Describe the bug**
When running the full suite on Kaggle, the 3 progress bars of "Ingesting batches" disappear after completion, but the "Computing single dataset checks" remains visible.

**To Reproduce**
View this notebook:
https://www.kaggle.com/code/itay94/notebook1b0b03205f


**Screenshots**
If applicable, add screenshots to help explain your problem.
![image](https://user-images.githubusercontent.com/20860465/161489710-ad6a065e-a7f9-4355-a10e-86b198c08928.png)

**Environment (please complete the following information):**
 - OS: Kaggle
 - Python Version: Kaggle kernel
 - Deepchecks Version: 0.6.0

**Additional context**
Add any other context about the problem here.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `deepchecks/vision/suite.py`
Content:
```
1 # ----------------------------------------------------------------------------
2 # Copyright (C) 2021-2022 Deepchecks (https://www.deepchecks.com)
3 #
4 # This file is part of Deepchecks.
5 # Deepchecks is distributed under the terms of the GNU Affero General
6 # Public License (version 3 or later).
7 # You should have received a copy of the GNU Affero General Public License
8 # along with Deepchecks.  If not, see <http://www.gnu.org/licenses/>.
9 # ----------------------------------------------------------------------------
10 #
11 """Module for base vision abstractions."""
12 # pylint: disable=broad-except,not-callable
13 import logging
14 from typing import Tuple, Mapping, Optional, Union, Dict, List
15 from collections import OrderedDict
16 
17 import torch
18 from torch import nn
19 from ignite.metrics import Metric
20 
21 from deepchecks.core.check_result import CheckFailure, CheckResult
22 from deepchecks.core.checks import DatasetKind
23 from deepchecks.core.suite import BaseSuite, SuiteResult
24 from deepchecks.core.display_suite import ProgressBar
25 from deepchecks.core.errors import DeepchecksNotSupportedError
26 from deepchecks.vision.base_checks import ModelOnlyCheck, SingleDatasetCheck, TrainTestCheck
27 from deepchecks.vision.context import Context
28 from deepchecks.vision.vision_data import VisionData
29 from deepchecks.vision.batch_wrapper import Batch
30 
31 
32 __all__ = ['Suite']
33 
34 
35 logger = logging.getLogger('deepchecks')
36 
37 
38 class Suite(BaseSuite):
39     """Tabular suite to run checks of types: TrainTestCheck, SingleDatasetCheck, ModelOnlyCheck."""
40 
41     @classmethod
42     def supported_checks(cls) -> Tuple:
43         """Return tuple of supported check types of this suite."""
44         return TrainTestCheck, SingleDatasetCheck, ModelOnlyCheck
45 
46     def run(
47             self,
48             train_dataset: Optional[VisionData] = None,
49             test_dataset: Optional[VisionData] = None,
50             model: nn.Module = None,
51             scorers: Mapping[str, Metric] = None,
52             scorers_per_class: Mapping[str, Metric] = None,
53             device: Union[str, torch.device, None] = 'cpu',
54             random_state: int = 42,
55             n_samples: Optional[int] = 10_000,
56     ) -> SuiteResult:
57         """Run all checks.
58 
59         Parameters
60         ----------
61         train_dataset: Optional[VisionData] , default None
62             object, representing data an estimator was fitted on
63         test_dataset : Optional[VisionData] , default None
64             object, representing data an estimator predicts on
65         model : nn.Module , default None
66             A scikit-learn-compatible fitted estimator instance
67         scorers : Mapping[str, Metric] , default None
68             dict of scorers names to scorer sklearn_name/function
69         scorers_per_class : Mapping[str, Metric], default None
70             dict of scorers for classification without averaging of the classes
71             See <a href=
72             "https://scikit-learn.org/stable/modules/model_evaluation.html#from-binary-to-multiclass-and-multilabel">
73             scikit-learn docs</a>
74         device : Union[str, torch.device], default: None
75             processing unit for use
76         random_state : int
77             A seed to set for pseudo-random functions
78         n_samples : int, default: 10,000
79             number of samples to draw from the dataset.
80         Returns
81         -------
82         SuiteResult
83             All results by all initialized checks
84         """
85         all_pbars = []
86         progress_bar = ProgressBar('Validating Input', 1, unit='')
87         all_pbars.append(progress_bar)
88         context = Context(
89             train_dataset,
90             test_dataset,
91             model,
92             scorers=scorers,
93             scorers_per_class=scorers_per_class,
94             device=device,
95             random_state=random_state,
96             n_samples=n_samples
97         )
98         progress_bar.inc_progress()
99 
100         results: Dict[
101             Union[str, int],
102             Union[CheckResult, CheckFailure]
103         ] = OrderedDict({})
104 
105         run_train_test_checks = train_dataset is not None and test_dataset is not None
106         non_single_checks = {k: check for k, check in self.checks.items() if not isinstance(check, SingleDatasetCheck)}
107 
108         # Initialize here all the checks that are not single dataset, since those are initialized inside the update loop
109         for index, check in non_single_checks.items():
110             try:
111                 check.initialize_run(context)
112             except Exception as exp:
113                 results[index] = CheckFailure(check, exp)
114 
115         if train_dataset is not None:
116             self._update_loop(
117                 context=context,
118                 run_train_test_checks=run_train_test_checks,
119                 results=results,
120                 dataset_kind=DatasetKind.TRAIN,
121                 progress_bars=all_pbars
122             )
123 
124         if test_dataset is not None:
125             self._update_loop(
126                 context=context,
127                 run_train_test_checks=run_train_test_checks,
128                 results=results,
129                 dataset_kind=DatasetKind.TEST,
130                 progress_bars=all_pbars
131             )
132 
133         # Need to compute only on not SingleDatasetCheck, since they computed inside the loop
134         if non_single_checks:
135             progress_bar = ProgressBar('Computing Checks', len(non_single_checks), unit='Check')
136             all_pbars.append(progress_bar)
137             for check_idx, check in non_single_checks.items():
138                 progress_bar.set_text(check.name())
139                 try:
140                     # if check index in results we had failure
141                     if check_idx not in results:
142                         result = check.finalize_check_result(check.compute(context))
143                         results[check_idx] = result
144                 except Exception as exp:
145                     results[check_idx] = CheckFailure(check, exp)
146                 progress_bar.inc_progress()
147 
148         # The results are ordered as they ran instead of in the order they were defined, therefore sort by key
149         sorted_result_values = [value for name, value in sorted(results.items(), key=lambda pair: str(pair[0]))]
150 
151         # Close all progress bars
152         for pbar in all_pbars:
153             pbar.close()
154 
155         footnote = context.get_is_sampled_footnote()
156         extra_info = [footnote] if footnote else []
157         return SuiteResult(self.name, sorted_result_values, extra_info)
158 
159     def _update_loop(
160         self,
161         context: Context,
162         run_train_test_checks: bool,
163         results: Dict[Union[str, int], Union[CheckResult, CheckFailure]],
164         dataset_kind: DatasetKind,
165         progress_bars: List
166     ):
167         type_suffix = ' - Test Dataset' if dataset_kind == DatasetKind.TEST else ' - Train Dataset'
168         vision_data = context.get_data_by_kind(dataset_kind)
169         n_batches = len(vision_data)
170         single_dataset_checks = {k: check for k, check in self.checks.items() if isinstance(check, SingleDatasetCheck)}
171 
172         # SingleDatasetChecks have different handling, need to initialize them here (to have them ready for different
173         # dataset kind)
174         for idx, check in single_dataset_checks.items():
175             try:
176                 check.initialize_run(context, dataset_kind=dataset_kind)
177             except Exception as exp:
178                 results[idx] = CheckFailure(check, exp, type_suffix)
179 
180         # Init cache of vision_data
181         vision_data.init_cache()
182 
183         progress_bar = ProgressBar('Ingesting Batches' + type_suffix, n_batches, unit='Batch')
184         progress_bars.append(progress_bar)
185 
186         # Run on all the batches
187         batch_start_index = 0
188         for batch_id, batch in enumerate(vision_data):
189             progress_bar.set_text(f'{100 * batch_id / (1. * n_batches):.0f}%')
190             batch = Batch(batch, context, dataset_kind, batch_start_index)
191             vision_data.update_cache(batch)
192             for check_idx, check in self.checks.items():
193                 # If index in results the check already failed before
194                 if check_idx in results:
195                     continue
196                 try:
197                     if isinstance(check, TrainTestCheck):
198                         if run_train_test_checks is True:
199                             check.update(context, batch, dataset_kind=dataset_kind)
200                         else:
201                             msg = 'Check is irrelevant if not supplied with both train and test datasets'
202                             results[check_idx] = self._get_unsupported_failure(check, msg)
203                     elif isinstance(check, SingleDatasetCheck):
204                         check.update(context, batch, dataset_kind=dataset_kind)
205                     elif isinstance(check, ModelOnlyCheck):
206                         pass
207                     else:
208                         raise TypeError(f'Don\'t know how to handle type {check.__class__.__name__} in suite.')
209                 except Exception as exp:
210                     results[check_idx] = CheckFailure(check, exp, type_suffix)
211 
212             batch_start_index += len(batch)
213             progress_bar.inc_progress()
214 
215         # SingleDatasetChecks have different handling. If we had failure in them need to add suffix to the index of
216         # the results, else need to compute it.
217         if single_dataset_checks:
218             progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix, len(single_dataset_checks),
219                                        unit='Check')
220             progress_bars.append(progress_bar)
221             for idx, check in single_dataset_checks.items():
222                 progress_bar.set_text(check.name())
223                 index_of_kind = str(idx) + type_suffix
224                 # If index in results we had a failure
225                 if idx in results:
226                     results[index_of_kind] = results.pop(idx)
227                     continue
228                 try:
229                     result = check.compute(context, dataset_kind=dataset_kind)
230                     result = check.finalize_check_result(result)
231                     # Update header with dataset type only if both train and test ran
232                     if run_train_test_checks:
233                         result.header = result.get_header() + type_suffix
234                     results[index_of_kind] = result
235                 except Exception as exp:
236                     results[index_of_kind] = CheckFailure(check, exp, type_suffix)
237                 progress_bar.inc_progress()
238 
239     @classmethod
240     def _get_unsupported_failure(cls, check, msg):
241         return CheckFailure(check, DeepchecksNotSupportedError(msg))
242
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/deepchecks/vision/suite.py b/deepchecks/vision/suite.py
--- a/deepchecks/vision/suite.py
+++ b/deepchecks/vision/suite.py
@@ -185,8 +185,7 @@
 
         # Run on all the batches
         batch_start_index = 0
-        for batch_id, batch in enumerate(vision_data):
-            progress_bar.set_text(f'{100 * batch_id / (1. * n_batches):.0f}%')
+        for batch in vision_data:
             batch = Batch(batch, context, dataset_kind, batch_start_index)
             vision_data.update_cache(batch)
             for check_idx, check in self.checks.items():
@@ -215,7 +214,8 @@
         # SingleDatasetChecks have different handling. If we had failure in them need to add suffix to the index of
         # the results, else need to compute it.
         if single_dataset_checks:
-            progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix, len(single_dataset_checks),
+            progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix,
+                                       len(single_dataset_checks),
                                        unit='Check')
             progress_bars.append(progress_bar)
             for idx, check in single_dataset_checks.items():
@@ -224,16 +224,16 @@
                 # If index in results we had a failure
                 if idx in results:
                     results[index_of_kind] = results.pop(idx)
-                    continue
-                try:
-                    result = check.compute(context, dataset_kind=dataset_kind)
-                    result = check.finalize_check_result(result)
-                    # Update header with dataset type only if both train and test ran
-                    if run_train_test_checks:
-                        result.header = result.get_header() + type_suffix
-                    results[index_of_kind] = result
-                except Exception as exp:
-                    results[index_of_kind] = CheckFailure(check, exp, type_suffix)
+                else:
+                    try:
+                        result = check.compute(context, dataset_kind=dataset_kind)
+                        result = check.finalize_check_result(result)
+                        # Update header with dataset type only if both train and test ran
+                        if run_train_test_checks:
+                            result.header = result.get_header() + type_suffix
+                        results[index_of_kind] = result
+                    except Exception as exp:
+                        results[index_of_kind] = CheckFailure(check, exp, type_suffix)
                 progress_bar.inc_progress()
 
     @classmethod

{"golden_diff": "diff --git a/deepchecks/vision/suite.py b/deepchecks/vision/suite.py\n--- a/deepchecks/vision/suite.py\n+++ b/deepchecks/vision/suite.py\n@@ -185,8 +185,7 @@\n \n         # Run on all the batches\n         batch_start_index = 0\n-        for batch_id, batch in enumerate(vision_data):\n-            progress_bar.set_text(f'{100 * batch_id / (1. * n_batches):.0f}%')\n+        for batch in vision_data:\n             batch = Batch(batch, context, dataset_kind, batch_start_index)\n             vision_data.update_cache(batch)\n             for check_idx, check in self.checks.items():\n@@ -215,7 +214,8 @@\n         # SingleDatasetChecks have different handling. If we had failure in them need to add suffix to the index of\n         # the results, else need to compute it.\n         if single_dataset_checks:\n-            progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix, len(single_dataset_checks),\n+            progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix,\n+                                       len(single_dataset_checks),\n                                        unit='Check')\n             progress_bars.append(progress_bar)\n             for idx, check in single_dataset_checks.items():\n@@ -224,16 +224,16 @@\n                 # If index in results we had a failure\n                 if idx in results:\n                     results[index_of_kind] = results.pop(idx)\n-                    continue\n-                try:\n-                    result = check.compute(context, dataset_kind=dataset_kind)\n-                    result = check.finalize_check_result(result)\n-                    # Update header with dataset type only if both train and test ran\n-                    if run_train_test_checks:\n-                        result.header = result.get_header() + type_suffix\n-                    results[index_of_kind] = result\n-                except Exception as exp:\n-                    results[index_of_kind] = CheckFailure(check, exp, type_suffix)\n+                else:\n+                    try:\n+                        result = check.compute(context, dataset_kind=dataset_kind)\n+                        result = check.finalize_check_result(result)\n+                        # Update header with dataset type only if both train and test ran\n+                        if run_train_test_checks:\n+                            result.header = result.get_header() + type_suffix\n+                        results[index_of_kind] = result\n+                    except Exception as exp:\n+                        results[index_of_kind] = CheckFailure(check, exp, type_suffix)\n                 progress_bar.inc_progress()\n \n     @classmethod\n", "issue": "[BUG] Widgeted progress bar remains visible after suite run\n**Describe the bug**\r\nWhen running the full suite on Kaggle, the 3 progress bars of \"Ingesting batches\" disappear after completion, but the \"Computing single dataset checks\" remains visible.\r\n\r\n**To Reproduce**\r\nView this notebook:\r\nhttps://www.kaggle.com/code/itay94/notebook1b0b03205f\r\n\r\n\r\n**Screenshots**\r\nIf applicable, add screenshots to help explain your problem.\r\n![image](https://user-images.githubusercontent.com/20860465/161489710-ad6a065e-a7f9-4355-a10e-86b198c08928.png)\r\n\r\n**Environment (please complete the following information):**\r\n - OS: Kaggle\r\n - Python Version: Kaggle kernel\r\n - Deepchecks Version: 0.6.0\r\n\r\n**Additional context**\r\nAdd any other context about the problem here.\r\n\n", "before_files": [{"content": "# ----------------------------------------------------------------------------\n# Copyright (C) 2021-2022 Deepchecks (https://www.deepchecks.com)\n#\n# This file is part of Deepchecks.\n# Deepchecks is distributed under the terms of the GNU Affero General\n# Public License (version 3 or later).\n# You should have received a copy of the GNU Affero General Public License\n# along with Deepchecks.  If not, see <http://www.gnu.org/licenses/>.\n# ----------------------------------------------------------------------------\n#\n\"\"\"Module for base vision abstractions.\"\"\"\n# pylint: disable=broad-except,not-callable\nimport logging\nfrom typing import Tuple, Mapping, Optional, Union, Dict, List\nfrom collections import OrderedDict\n\nimport torch\nfrom torch import nn\nfrom ignite.metrics import Metric\n\nfrom deepchecks.core.check_result import CheckFailure, CheckResult\nfrom deepchecks.core.checks import DatasetKind\nfrom deepchecks.core.suite import BaseSuite, SuiteResult\nfrom deepchecks.core.display_suite import ProgressBar\nfrom deepchecks.core.errors import DeepchecksNotSupportedError\nfrom deepchecks.vision.base_checks import ModelOnlyCheck, SingleDatasetCheck, TrainTestCheck\nfrom deepchecks.vision.context import Context\nfrom deepchecks.vision.vision_data import VisionData\nfrom deepchecks.vision.batch_wrapper import Batch\n\n\n__all__ = ['Suite']\n\n\nlogger = logging.getLogger('deepchecks')\n\n\nclass Suite(BaseSuite):\n    \"\"\"Tabular suite to run checks of types: TrainTestCheck, SingleDatasetCheck, ModelOnlyCheck.\"\"\"\n\n    @classmethod\n    def supported_checks(cls) -> Tuple:\n        \"\"\"Return tuple of supported check types of this suite.\"\"\"\n        return TrainTestCheck, SingleDatasetCheck, ModelOnlyCheck\n\n    def run(\n            self,\n            train_dataset: Optional[VisionData] = None,\n            test_dataset: Optional[VisionData] = None,\n            model: nn.Module = None,\n            scorers: Mapping[str, Metric] = None,\n            scorers_per_class: Mapping[str, Metric] = None,\n            device: Union[str, torch.device, None] = 'cpu',\n            random_state: int = 42,\n            n_samples: Optional[int] = 10_000,\n    ) -> SuiteResult:\n        \"\"\"Run all checks.\n\n        Parameters\n        ----------\n        train_dataset: Optional[VisionData] , default None\n            object, representing data an estimator was fitted on\n        test_dataset : Optional[VisionData] , default None\n            object, representing data an estimator predicts on\n        model : nn.Module , default None\n            A scikit-learn-compatible fitted estimator instance\n        scorers : Mapping[str, Metric] , default None\n            dict of scorers names to scorer sklearn_name/function\n        scorers_per_class : Mapping[str, Metric], default None\n            dict of scorers for classification without averaging of the classes\n            See <a href=\n            \"https://scikit-learn.org/stable/modules/model_evaluation.html#from-binary-to-multiclass-and-multilabel\">\n            scikit-learn docs</a>\n        device : Union[str, torch.device], default: None\n            processing unit for use\n        random_state : int\n            A seed to set for pseudo-random functions\n        n_samples : int, default: 10,000\n            number of samples to draw from the dataset.\n        Returns\n        -------\n        SuiteResult\n            All results by all initialized checks\n        \"\"\"\n        all_pbars = []\n        progress_bar = ProgressBar('Validating Input', 1, unit='')\n        all_pbars.append(progress_bar)\n        context = Context(\n            train_dataset,\n            test_dataset,\n            model,\n            scorers=scorers,\n            scorers_per_class=scorers_per_class,\n            device=device,\n            random_state=random_state,\n            n_samples=n_samples\n        )\n        progress_bar.inc_progress()\n\n        results: Dict[\n            Union[str, int],\n            Union[CheckResult, CheckFailure]\n        ] = OrderedDict({})\n\n        run_train_test_checks = train_dataset is not None and test_dataset is not None\n        non_single_checks = {k: check for k, check in self.checks.items() if not isinstance(check, SingleDatasetCheck)}\n\n        # Initialize here all the checks that are not single dataset, since those are initialized inside the update loop\n        for index, check in non_single_checks.items():\n            try:\n                check.initialize_run(context)\n            except Exception as exp:\n                results[index] = CheckFailure(check, exp)\n\n        if train_dataset is not None:\n            self._update_loop(\n                context=context,\n                run_train_test_checks=run_train_test_checks,\n                results=results,\n                dataset_kind=DatasetKind.TRAIN,\n                progress_bars=all_pbars\n            )\n\n        if test_dataset is not None:\n            self._update_loop(\n                context=context,\n                run_train_test_checks=run_train_test_checks,\n                results=results,\n                dataset_kind=DatasetKind.TEST,\n                progress_bars=all_pbars\n            )\n\n        # Need to compute only on not SingleDatasetCheck, since they computed inside the loop\n        if non_single_checks:\n            progress_bar = ProgressBar('Computing Checks', len(non_single_checks), unit='Check')\n            all_pbars.append(progress_bar)\n            for check_idx, check in non_single_checks.items():\n                progress_bar.set_text(check.name())\n                try:\n                    # if check index in results we had failure\n                    if check_idx not in results:\n                        result = check.finalize_check_result(check.compute(context))\n                        results[check_idx] = result\n                except Exception as exp:\n                    results[check_idx] = CheckFailure(check, exp)\n                progress_bar.inc_progress()\n\n        # The results are ordered as they ran instead of in the order they were defined, therefore sort by key\n        sorted_result_values = [value for name, value in sorted(results.items(), key=lambda pair: str(pair[0]))]\n\n        # Close all progress bars\n        for pbar in all_pbars:\n            pbar.close()\n\n        footnote = context.get_is_sampled_footnote()\n        extra_info = [footnote] if footnote else []\n        return SuiteResult(self.name, sorted_result_values, extra_info)\n\n    def _update_loop(\n        self,\n        context: Context,\n        run_train_test_checks: bool,\n        results: Dict[Union[str, int], Union[CheckResult, CheckFailure]],\n        dataset_kind: DatasetKind,\n        progress_bars: List\n    ):\n        type_suffix = ' - Test Dataset' if dataset_kind == DatasetKind.TEST else ' - Train Dataset'\n        vision_data = context.get_data_by_kind(dataset_kind)\n        n_batches = len(vision_data)\n        single_dataset_checks = {k: check for k, check in self.checks.items() if isinstance(check, SingleDatasetCheck)}\n\n        # SingleDatasetChecks have different handling, need to initialize them here (to have them ready for different\n        # dataset kind)\n        for idx, check in single_dataset_checks.items():\n            try:\n                check.initialize_run(context, dataset_kind=dataset_kind)\n            except Exception as exp:\n                results[idx] = CheckFailure(check, exp, type_suffix)\n\n        # Init cache of vision_data\n        vision_data.init_cache()\n\n        progress_bar = ProgressBar('Ingesting Batches' + type_suffix, n_batches, unit='Batch')\n        progress_bars.append(progress_bar)\n\n        # Run on all the batches\n        batch_start_index = 0\n        for batch_id, batch in enumerate(vision_data):\n            progress_bar.set_text(f'{100 * batch_id / (1. * n_batches):.0f}%')\n            batch = Batch(batch, context, dataset_kind, batch_start_index)\n            vision_data.update_cache(batch)\n            for check_idx, check in self.checks.items():\n                # If index in results the check already failed before\n                if check_idx in results:\n                    continue\n                try:\n                    if isinstance(check, TrainTestCheck):\n                        if run_train_test_checks is True:\n                            check.update(context, batch, dataset_kind=dataset_kind)\n                        else:\n                            msg = 'Check is irrelevant if not supplied with both train and test datasets'\n                            results[check_idx] = self._get_unsupported_failure(check, msg)\n                    elif isinstance(check, SingleDatasetCheck):\n                        check.update(context, batch, dataset_kind=dataset_kind)\n                    elif isinstance(check, ModelOnlyCheck):\n                        pass\n                    else:\n                        raise TypeError(f'Don\\'t know how to handle type {check.__class__.__name__} in suite.')\n                except Exception as exp:\n                    results[check_idx] = CheckFailure(check, exp, type_suffix)\n\n            batch_start_index += len(batch)\n            progress_bar.inc_progress()\n\n        # SingleDatasetChecks have different handling. If we had failure in them need to add suffix to the index of\n        # the results, else need to compute it.\n        if single_dataset_checks:\n            progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix, len(single_dataset_checks),\n                                       unit='Check')\n            progress_bars.append(progress_bar)\n            for idx, check in single_dataset_checks.items():\n                progress_bar.set_text(check.name())\n                index_of_kind = str(idx) + type_suffix\n                # If index in results we had a failure\n                if idx in results:\n                    results[index_of_kind] = results.pop(idx)\n                    continue\n                try:\n                    result = check.compute(context, dataset_kind=dataset_kind)\n                    result = check.finalize_check_result(result)\n                    # Update header with dataset type only if both train and test ran\n                    if run_train_test_checks:\n                        result.header = result.get_header() + type_suffix\n                    results[index_of_kind] = result\n                except Exception as exp:\n                    results[index_of_kind] = CheckFailure(check, exp, type_suffix)\n                progress_bar.inc_progress()\n\n    @classmethod\n    def _get_unsupported_failure(cls, check, msg):\n        return CheckFailure(check, DeepchecksNotSupportedError(msg))\n", "path": "deepchecks/vision/suite.py"}], "after_files": [{"content": "# ----------------------------------------------------------------------------\n# Copyright (C) 2021-2022 Deepchecks (https://www.deepchecks.com)\n#\n# This file is part of Deepchecks.\n# Deepchecks is distributed under the terms of the GNU Affero General\n# Public License (version 3 or later).\n# You should have received a copy of the GNU Affero General Public License\n# along with Deepchecks.  If not, see <http://www.gnu.org/licenses/>.\n# ----------------------------------------------------------------------------\n#\n\"\"\"Module for base vision abstractions.\"\"\"\n# pylint: disable=broad-except,not-callable\nimport logging\nfrom typing import Tuple, Mapping, Optional, Union, Dict, List\nfrom collections import OrderedDict\n\nimport torch\nfrom torch import nn\nfrom ignite.metrics import Metric\n\nfrom deepchecks.core.check_result import CheckFailure, CheckResult\nfrom deepchecks.core.checks import DatasetKind\nfrom deepchecks.core.suite import BaseSuite, SuiteResult\nfrom deepchecks.core.display_suite import ProgressBar\nfrom deepchecks.core.errors import DeepchecksNotSupportedError\nfrom deepchecks.vision.base_checks import ModelOnlyCheck, SingleDatasetCheck, TrainTestCheck\nfrom deepchecks.vision.context import Context\nfrom deepchecks.vision.vision_data import VisionData\nfrom deepchecks.vision.batch_wrapper import Batch\n\n\n__all__ = ['Suite']\n\n\nlogger = logging.getLogger('deepchecks')\n\n\nclass Suite(BaseSuite):\n    \"\"\"Tabular suite to run checks of types: TrainTestCheck, SingleDatasetCheck, ModelOnlyCheck.\"\"\"\n\n    @classmethod\n    def supported_checks(cls) -> Tuple:\n        \"\"\"Return tuple of supported check types of this suite.\"\"\"\n        return TrainTestCheck, SingleDatasetCheck, ModelOnlyCheck\n\n    def run(\n            self,\n            train_dataset: Optional[VisionData] = None,\n            test_dataset: Optional[VisionData] = None,\n            model: nn.Module = None,\n            scorers: Mapping[str, Metric] = None,\n            scorers_per_class: Mapping[str, Metric] = None,\n            device: Union[str, torch.device, None] = 'cpu',\n            random_state: int = 42,\n            n_samples: Optional[int] = 10_000,\n    ) -> SuiteResult:\n        \"\"\"Run all checks.\n\n        Parameters\n        ----------\n        train_dataset: Optional[VisionData] , default None\n            object, representing data an estimator was fitted on\n        test_dataset : Optional[VisionData] , default None\n            object, representing data an estimator predicts on\n        model : nn.Module , default None\n            A scikit-learn-compatible fitted estimator instance\n        scorers : Mapping[str, Metric] , default None\n            dict of scorers names to scorer sklearn_name/function\n        scorers_per_class : Mapping[str, Metric], default None\n            dict of scorers for classification without averaging of the classes\n            See <a href=\n            \"https://scikit-learn.org/stable/modules/model_evaluation.html#from-binary-to-multiclass-and-multilabel\">\n            scikit-learn docs</a>\n        device : Union[str, torch.device], default: None\n            processing unit for use\n        random_state : int\n            A seed to set for pseudo-random functions\n        n_samples : int, default: 10,000\n            number of samples to draw from the dataset.\n        Returns\n        -------\n        SuiteResult\n            All results by all initialized checks\n        \"\"\"\n        all_pbars = []\n        progress_bar = ProgressBar('Validating Input', 1, unit='')\n        all_pbars.append(progress_bar)\n        context = Context(\n            train_dataset,\n            test_dataset,\n            model,\n            scorers=scorers,\n            scorers_per_class=scorers_per_class,\n            device=device,\n            random_state=random_state,\n            n_samples=n_samples\n        )\n        progress_bar.inc_progress()\n\n        results: Dict[\n            Union[str, int],\n            Union[CheckResult, CheckFailure]\n        ] = OrderedDict({})\n\n        run_train_test_checks = train_dataset is not None and test_dataset is not None\n        non_single_checks = {k: check for k, check in self.checks.items() if not isinstance(check, SingleDatasetCheck)}\n\n        # Initialize here all the checks that are not single dataset, since those are initialized inside the update loop\n        for index, check in non_single_checks.items():\n            try:\n                check.initialize_run(context)\n            except Exception as exp:\n                results[index] = CheckFailure(check, exp)\n\n        if train_dataset is not None:\n            self._update_loop(\n                context=context,\n                run_train_test_checks=run_train_test_checks,\n                results=results,\n                dataset_kind=DatasetKind.TRAIN,\n                progress_bars=all_pbars\n            )\n\n        if test_dataset is not None:\n            self._update_loop(\n                context=context,\n                run_train_test_checks=run_train_test_checks,\n                results=results,\n                dataset_kind=DatasetKind.TEST,\n                progress_bars=all_pbars\n            )\n\n        # Need to compute only on not SingleDatasetCheck, since they computed inside the loop\n        if non_single_checks:\n            progress_bar = ProgressBar('Computing Checks', len(non_single_checks), unit='Check')\n            all_pbars.append(progress_bar)\n            for check_idx, check in non_single_checks.items():\n                progress_bar.set_text(check.name())\n                try:\n                    # if check index in results we had failure\n                    if check_idx not in results:\n                        result = check.finalize_check_result(check.compute(context))\n                        results[check_idx] = result\n                except Exception as exp:\n                    results[check_idx] = CheckFailure(check, exp)\n                progress_bar.inc_progress()\n\n        # The results are ordered as they ran instead of in the order they were defined, therefore sort by key\n        sorted_result_values = [value for name, value in sorted(results.items(), key=lambda pair: str(pair[0]))]\n\n        # Close all progress bars\n        for pbar in all_pbars:\n            pbar.close()\n\n        footnote = context.get_is_sampled_footnote()\n        extra_info = [footnote] if footnote else []\n        return SuiteResult(self.name, sorted_result_values, extra_info)\n\n    def _update_loop(\n        self,\n        context: Context,\n        run_train_test_checks: bool,\n        results: Dict[Union[str, int], Union[CheckResult, CheckFailure]],\n        dataset_kind: DatasetKind,\n        progress_bars: List\n    ):\n        type_suffix = ' - Test Dataset' if dataset_kind == DatasetKind.TEST else ' - Train Dataset'\n        vision_data = context.get_data_by_kind(dataset_kind)\n        n_batches = len(vision_data)\n        single_dataset_checks = {k: check for k, check in self.checks.items() if isinstance(check, SingleDatasetCheck)}\n\n        # SingleDatasetChecks have different handling, need to initialize them here (to have them ready for different\n        # dataset kind)\n        for idx, check in single_dataset_checks.items():\n            try:\n                check.initialize_run(context, dataset_kind=dataset_kind)\n            except Exception as exp:\n                results[idx] = CheckFailure(check, exp, type_suffix)\n\n        # Init cache of vision_data\n        vision_data.init_cache()\n\n        progress_bar = ProgressBar('Ingesting Batches' + type_suffix, n_batches, unit='Batch')\n        progress_bars.append(progress_bar)\n\n        # Run on all the batches\n        batch_start_index = 0\n        for batch in vision_data:\n            batch = Batch(batch, context, dataset_kind, batch_start_index)\n            vision_data.update_cache(batch)\n            for check_idx, check in self.checks.items():\n                # If index in results the check already failed before\n                if check_idx in results:\n                    continue\n                try:\n                    if isinstance(check, TrainTestCheck):\n                        if run_train_test_checks is True:\n                            check.update(context, batch, dataset_kind=dataset_kind)\n                        else:\n                            msg = 'Check is irrelevant if not supplied with both train and test datasets'\n                            results[check_idx] = self._get_unsupported_failure(check, msg)\n                    elif isinstance(check, SingleDatasetCheck):\n                        check.update(context, batch, dataset_kind=dataset_kind)\n                    elif isinstance(check, ModelOnlyCheck):\n                        pass\n                    else:\n                        raise TypeError(f'Don\\'t know how to handle type {check.__class__.__name__} in suite.')\n                except Exception as exp:\n                    results[check_idx] = CheckFailure(check, exp, type_suffix)\n\n            batch_start_index += len(batch)\n            progress_bar.inc_progress()\n\n        # SingleDatasetChecks have different handling. If we had failure in them need to add suffix to the index of\n        # the results, else need to compute it.\n        if single_dataset_checks:\n            progress_bar = ProgressBar('Computing Single Dataset Checks' + type_suffix,\n                                       len(single_dataset_checks),\n                                       unit='Check')\n            progress_bars.append(progress_bar)\n            for idx, check in single_dataset_checks.items():\n                progress_bar.set_text(check.name())\n                index_of_kind = str(idx) + type_suffix\n                # If index in results we had a failure\n                if idx in results:\n                    results[index_of_kind] = results.pop(idx)\n                else:\n                    try:\n                        result = check.compute(context, dataset_kind=dataset_kind)\n                        result = check.finalize_check_result(result)\n                        # Update header with dataset type only if both train and test ran\n                        if run_train_test_checks:\n                            result.header = result.get_header() + type_suffix\n                        results[index_of_kind] = result\n                    except Exception as exp:\n                        results[index_of_kind] = CheckFailure(check, exp, type_suffix)\n                progress_bar.inc_progress()\n\n    @classmethod\n    def _get_unsupported_failure(cls, check, msg):\n        return CheckFailure(check, DeepchecksNotSupportedError(msg))\n", "path": "deepchecks/vision/suite.py"}]}