We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
SustainBenchCropYield download doesn't work
### Description

Downloading the SustainBenchCropYield dataset doesn't work as expected

### Steps to reproduce

```
ds = SustainBenchCropYield("data/", download=True)
```

This downloads a file called `soybeans` then fails unzipping `soybeans.zip`. Works if you rename to .zip and unzip manually.

### Version

0.6.0.dev0

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchgeo/datasets/sustainbench_crop_yield.py`
Content:
```
1 # Copyright (c) Microsoft Corporation. All rights reserved.
2 # Licensed under the MIT License.
3 
4 """SustainBench Crop Yield dataset."""
5 
6 import os
7 from typing import Any, Callable, Optional
8 
9 import matplotlib.pyplot as plt
10 import numpy as np
11 import torch
12 from matplotlib.figure import Figure
13 from torch import Tensor
14 
15 from .geo import NonGeoDataset
16 from .utils import DatasetNotFoundError, download_url, extract_archive
17 
18 
19 class SustainBenchCropYield(NonGeoDataset):
20     """SustainBench Crop Yield Dataset.
21 
22     This dataset contains MODIS band histograms and soybean yield
23     estimates for selected counties in the USA, Argentina and Brazil.
24     The dataset is part of the
25     `SustainBench <https://sustainlab-group.github.io/sustainbench/docs/datasets/sdg2/crop_yield.html>`_
26     datasets for tackling the UN Sustainable Development Goals (SDGs).
27 
28     Dataset Format:
29 
30     * .npz files of stacked samples
31 
32     Dataset Features:
33 
34     * input histogram of 7 surface reflectance and 2 surface temperature
35       bands from MODIS pixel values in 32 ranges across 32 timesteps
36       resulting in 32x32x9 input images
37     * regression target value of soybean yield in metric tonnes per
38       harvested hectare
39 
40     If you use this dataset in your research, please cite:
41 
42     * https://doi.org/10.1145/3209811.3212707
43     * https://doi.org/10.1609/aaai.v31i1.11172
44 
45     .. versionadded:: 0.5
46     """  # noqa: E501
47 
48     valid_countries = ["usa", "brazil", "argentina"]
49 
50     md5 = "c2794e59512c897d9bea77b112848122"
51 
52     url = "https://drive.google.com/file/d/1odwkI1hiE5rMZ4VfM0hOXzlFR4NbhrfU/view?usp=share_link"  # noqa: E501
53 
54     dir = "soybeans"
55 
56     valid_splits = ["train", "dev", "test"]
57 
58     def __init__(
59         self,
60         root: str = "data",
61         split: str = "train",
62         countries: list[str] = ["usa"],
63         transforms: Optional[Callable[[dict[str, Any]], dict[str, Any]]] = None,
64         download: bool = False,
65         checksum: bool = False,
66     ) -> None:
67         """Initialize a new Dataset instance.
68 
69         Args:
70             root: root directory where dataset can be found
71             split: one of "train", "dev", or "test"
72             countries: which countries to include in the dataset
73             transforms: a function/transform that takes an input sample
74                 and returns a transformed version
75             download: if True, download dataset and store it in the root directory
76             checksum: if True, check the MD5 after downloading files (may be slow)
77 
78         Raises:
79             AssertionError: if ``countries`` contains invalid countries or if ``split``
80                 is invalid
81             DatasetNotFoundError: If dataset is not found and *download* is False.
82         """
83         assert set(countries).issubset(
84             self.valid_countries
85         ), f"Please choose a subset of these valid countried: {self.valid_countries}."
86         self.countries = countries
87 
88         assert (
89             split in self.valid_splits
90         ), f"Pleas choose one of these valid data splits {self.valid_splits}."
91         self.split = split
92 
93         self.root = root
94         self.transforms = transforms
95         self.download = download
96         self.checksum = checksum
97 
98         self._verify()
99         self.collection = self.retrieve_collection()
100 
101     def __len__(self) -> int:
102         """Return the number of data points in the dataset.
103 
104         Returns:
105             length of the dataset
106         """
107         return len(self.collection)
108 
109     def __getitem__(self, index: int) -> dict[str, Tensor]:
110         """Return an index within the dataset.
111 
112         Args:
113             index: index to return
114 
115         Returns:
116             data and label at that index
117         """
118         input_file_path, sample_idx = self.collection[index]
119 
120         sample: dict[str, Tensor] = {
121             "image": self._load_image(input_file_path, sample_idx)
122         }
123         sample.update(self._load_features(input_file_path, sample_idx))
124 
125         if self.transforms is not None:
126             sample = self.transforms(sample)
127 
128         return sample
129 
130     def _load_image(self, path: str, sample_idx: int) -> Tensor:
131         """Load input image.
132 
133         Args:
134             path: path to input npz collection
135             sample_idx: what sample to index from the npz collection
136 
137         Returns:
138             input image as tensor
139         """
140         arr = np.load(path)["data"][sample_idx]
141         # return [channel, height, width]
142         return torch.from_numpy(arr).permute(2, 0, 1).to(torch.float32)
143 
144     def _load_features(self, path: str, sample_idx: int) -> dict[str, Tensor]:
145         """Load features value.
146 
147         Args:
148             path: path to image npz collection
149             sample_idx: what sample to index from the npz collection
150 
151         Returns:
152             target regression value
153         """
154         target_file_path = path.replace("_hists", "_yields")
155         target = np.load(target_file_path)["data"][sample_idx]
156 
157         years_file_path = path.replace("_hists", "_years")
158         year = int(np.load(years_file_path)["data"][sample_idx])
159 
160         ndvi_file_path = path.replace("_hists", "_ndvi")
161         ndvi = np.load(ndvi_file_path)["data"][sample_idx]
162 
163         features = {
164             "label": torch.tensor(target).to(torch.float32),
165             "year": torch.tensor(year),
166             "ndvi": torch.from_numpy(ndvi).to(dtype=torch.float32),
167         }
168         return features
169 
170     def retrieve_collection(self) -> list[tuple[str, int]]:
171         """Retrieve the collection.
172 
173         Returns:
174             path and index to dataset samples
175         """
176         collection = []
177         for country in self.countries:
178             file_path = os.path.join(
179                 self.root, self.dir, country, f"{self.split}_hists.npz"
180             )
181             npz_file = np.load(file_path)
182             num_data_points = npz_file["data"].shape[0]
183             for idx in range(num_data_points):
184                 collection.append((file_path, idx))
185 
186         return collection
187 
188     def _verify(self) -> None:
189         """Verify the integrity of the dataset."""
190         # Check if the extracted files already exist
191         pathname = os.path.join(self.root, self.dir)
192         if os.path.exists(pathname):
193             return
194 
195         # Check if the zip files have already been downloaded
196         pathname = os.path.join(self.root, self.dir) + ".zip"
197         if os.path.exists(pathname):
198             self._extract()
199             return
200 
201         # Check if the user requested to download the dataset
202         if not self.download:
203             raise DatasetNotFoundError(self)
204 
205         # Download the dataset
206         self._download()
207         self._extract()
208 
209     def _download(self) -> None:
210         """Download the dataset and extract it."""
211         download_url(
212             self.url,
213             self.root,
214             filename=self.dir,
215             md5=self.md5 if self.checksum else None,
216         )
217         self._extract()
218 
219     def _extract(self) -> None:
220         """Extract the dataset."""
221         zipfile_path = os.path.join(self.root, self.dir) + ".zip"
222         extract_archive(zipfile_path, self.root)
223 
224     def plot(
225         self,
226         sample: dict[str, Any],
227         band_idx: int = 0,
228         show_titles: bool = True,
229         suptitle: Optional[str] = None,
230     ) -> Figure:
231         """Plot a sample from the dataset.
232 
233         Args:
234             sample: a sample return by :meth:`__getitem__`
235             band_idx: which of the nine histograms to index
236             show_titles: flag indicating whether to show titles above each panel
237             suptitle: optional suptitle to use for figure
238 
239         Returns:
240             a matplotlib Figure with the rendered sample
241 
242         """
243         image, label = sample["image"], sample["label"].item()
244 
245         showing_predictions = "prediction" in sample
246         if showing_predictions:
247             prediction = sample["prediction"].item()
248 
249         fig, ax = plt.subplots(1, 1, figsize=(10, 10))
250 
251         ax.imshow(image.permute(1, 2, 0)[:, :, band_idx])
252         ax.axis("off")
253 
254         if show_titles:
255             title = f"Label: {label:.3f}"
256             if showing_predictions:
257                 title += f"\nPrediction: {prediction:.3f}"
258             ax.set_title(title)
259 
260         if suptitle is not None:
261             plt.suptitle(suptitle)
262 
263         return fig
264
```
--- 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/torchgeo/datasets/sustainbench_crop_yield.py b/torchgeo/datasets/sustainbench_crop_yield.py
--- a/torchgeo/datasets/sustainbench_crop_yield.py
+++ b/torchgeo/datasets/sustainbench_crop_yield.py
@@ -211,7 +211,7 @@
         download_url(
             self.url,
             self.root,
-            filename=self.dir,
+            filename=self.dir + ".zip",
             md5=self.md5 if self.checksum else None,
         )
         self._extract()

{"golden_diff": "diff --git a/torchgeo/datasets/sustainbench_crop_yield.py b/torchgeo/datasets/sustainbench_crop_yield.py\n--- a/torchgeo/datasets/sustainbench_crop_yield.py\n+++ b/torchgeo/datasets/sustainbench_crop_yield.py\n@@ -211,7 +211,7 @@\n         download_url(\n             self.url,\n             self.root,\n-            filename=self.dir,\n+            filename=self.dir + \".zip\",\n             md5=self.md5 if self.checksum else None,\n         )\n         self._extract()\n", "issue": "SustainBenchCropYield download doesn't work\n### Description\n\nDownloading the SustainBenchCropYield dataset doesn't work as expected\n\n### Steps to reproduce\n\n```\r\nds = SustainBenchCropYield(\"data/\", download=True)\r\n```\r\n\r\nThis downloads a file called `soybeans` then fails unzipping `soybeans.zip`. Works if you rename to .zip and unzip manually.\n\n### Version\n\n0.6.0.dev0\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"SustainBench Crop Yield dataset.\"\"\"\n\nimport os\nfrom typing import Any, Callable, Optional\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom matplotlib.figure import Figure\nfrom torch import Tensor\n\nfrom .geo import NonGeoDataset\nfrom .utils import DatasetNotFoundError, download_url, extract_archive\n\n\nclass SustainBenchCropYield(NonGeoDataset):\n    \"\"\"SustainBench Crop Yield Dataset.\n\n    This dataset contains MODIS band histograms and soybean yield\n    estimates for selected counties in the USA, Argentina and Brazil.\n    The dataset is part of the\n    `SustainBench <https://sustainlab-group.github.io/sustainbench/docs/datasets/sdg2/crop_yield.html>`_\n    datasets for tackling the UN Sustainable Development Goals (SDGs).\n\n    Dataset Format:\n\n    * .npz files of stacked samples\n\n    Dataset Features:\n\n    * input histogram of 7 surface reflectance and 2 surface temperature\n      bands from MODIS pixel values in 32 ranges across 32 timesteps\n      resulting in 32x32x9 input images\n    * regression target value of soybean yield in metric tonnes per\n      harvested hectare\n\n    If you use this dataset in your research, please cite:\n\n    * https://doi.org/10.1145/3209811.3212707\n    * https://doi.org/10.1609/aaai.v31i1.11172\n\n    .. versionadded:: 0.5\n    \"\"\"  # noqa: E501\n\n    valid_countries = [\"usa\", \"brazil\", \"argentina\"]\n\n    md5 = \"c2794e59512c897d9bea77b112848122\"\n\n    url = \"https://drive.google.com/file/d/1odwkI1hiE5rMZ4VfM0hOXzlFR4NbhrfU/view?usp=share_link\"  # noqa: E501\n\n    dir = \"soybeans\"\n\n    valid_splits = [\"train\", \"dev\", \"test\"]\n\n    def __init__(\n        self,\n        root: str = \"data\",\n        split: str = \"train\",\n        countries: list[str] = [\"usa\"],\n        transforms: Optional[Callable[[dict[str, Any]], dict[str, Any]]] = None,\n        download: bool = False,\n        checksum: bool = False,\n    ) -> None:\n        \"\"\"Initialize a new Dataset instance.\n\n        Args:\n            root: root directory where dataset can be found\n            split: one of \"train\", \"dev\", or \"test\"\n            countries: which countries to include in the dataset\n            transforms: a function/transform that takes an input sample\n                and returns a transformed version\n            download: if True, download dataset and store it in the root directory\n            checksum: if True, check the MD5 after downloading files (may be slow)\n\n        Raises:\n            AssertionError: if ``countries`` contains invalid countries or if ``split``\n                is invalid\n            DatasetNotFoundError: If dataset is not found and *download* is False.\n        \"\"\"\n        assert set(countries).issubset(\n            self.valid_countries\n        ), f\"Please choose a subset of these valid countried: {self.valid_countries}.\"\n        self.countries = countries\n\n        assert (\n            split in self.valid_splits\n        ), f\"Pleas choose one of these valid data splits {self.valid_splits}.\"\n        self.split = split\n\n        self.root = root\n        self.transforms = transforms\n        self.download = download\n        self.checksum = checksum\n\n        self._verify()\n        self.collection = self.retrieve_collection()\n\n    def __len__(self) -> int:\n        \"\"\"Return the number of data points in the dataset.\n\n        Returns:\n            length of the dataset\n        \"\"\"\n        return len(self.collection)\n\n    def __getitem__(self, index: int) -> dict[str, Tensor]:\n        \"\"\"Return an index within the dataset.\n\n        Args:\n            index: index to return\n\n        Returns:\n            data and label at that index\n        \"\"\"\n        input_file_path, sample_idx = self.collection[index]\n\n        sample: dict[str, Tensor] = {\n            \"image\": self._load_image(input_file_path, sample_idx)\n        }\n        sample.update(self._load_features(input_file_path, sample_idx))\n\n        if self.transforms is not None:\n            sample = self.transforms(sample)\n\n        return sample\n\n    def _load_image(self, path: str, sample_idx: int) -> Tensor:\n        \"\"\"Load input image.\n\n        Args:\n            path: path to input npz collection\n            sample_idx: what sample to index from the npz collection\n\n        Returns:\n            input image as tensor\n        \"\"\"\n        arr = np.load(path)[\"data\"][sample_idx]\n        # return [channel, height, width]\n        return torch.from_numpy(arr).permute(2, 0, 1).to(torch.float32)\n\n    def _load_features(self, path: str, sample_idx: int) -> dict[str, Tensor]:\n        \"\"\"Load features value.\n\n        Args:\n            path: path to image npz collection\n            sample_idx: what sample to index from the npz collection\n\n        Returns:\n            target regression value\n        \"\"\"\n        target_file_path = path.replace(\"_hists\", \"_yields\")\n        target = np.load(target_file_path)[\"data\"][sample_idx]\n\n        years_file_path = path.replace(\"_hists\", \"_years\")\n        year = int(np.load(years_file_path)[\"data\"][sample_idx])\n\n        ndvi_file_path = path.replace(\"_hists\", \"_ndvi\")\n        ndvi = np.load(ndvi_file_path)[\"data\"][sample_idx]\n\n        features = {\n            \"label\": torch.tensor(target).to(torch.float32),\n            \"year\": torch.tensor(year),\n            \"ndvi\": torch.from_numpy(ndvi).to(dtype=torch.float32),\n        }\n        return features\n\n    def retrieve_collection(self) -> list[tuple[str, int]]:\n        \"\"\"Retrieve the collection.\n\n        Returns:\n            path and index to dataset samples\n        \"\"\"\n        collection = []\n        for country in self.countries:\n            file_path = os.path.join(\n                self.root, self.dir, country, f\"{self.split}_hists.npz\"\n            )\n            npz_file = np.load(file_path)\n            num_data_points = npz_file[\"data\"].shape[0]\n            for idx in range(num_data_points):\n                collection.append((file_path, idx))\n\n        return collection\n\n    def _verify(self) -> None:\n        \"\"\"Verify the integrity of the dataset.\"\"\"\n        # Check if the extracted files already exist\n        pathname = os.path.join(self.root, self.dir)\n        if os.path.exists(pathname):\n            return\n\n        # Check if the zip files have already been downloaded\n        pathname = os.path.join(self.root, self.dir) + \".zip\"\n        if os.path.exists(pathname):\n            self._extract()\n            return\n\n        # Check if the user requested to download the dataset\n        if not self.download:\n            raise DatasetNotFoundError(self)\n\n        # Download the dataset\n        self._download()\n        self._extract()\n\n    def _download(self) -> None:\n        \"\"\"Download the dataset and extract it.\"\"\"\n        download_url(\n            self.url,\n            self.root,\n            filename=self.dir,\n            md5=self.md5 if self.checksum else None,\n        )\n        self._extract()\n\n    def _extract(self) -> None:\n        \"\"\"Extract the dataset.\"\"\"\n        zipfile_path = os.path.join(self.root, self.dir) + \".zip\"\n        extract_archive(zipfile_path, self.root)\n\n    def plot(\n        self,\n        sample: dict[str, Any],\n        band_idx: int = 0,\n        show_titles: bool = True,\n        suptitle: Optional[str] = None,\n    ) -> Figure:\n        \"\"\"Plot a sample from the dataset.\n\n        Args:\n            sample: a sample return by :meth:`__getitem__`\n            band_idx: which of the nine histograms to index\n            show_titles: flag indicating whether to show titles above each panel\n            suptitle: optional suptitle to use for figure\n\n        Returns:\n            a matplotlib Figure with the rendered sample\n\n        \"\"\"\n        image, label = sample[\"image\"], sample[\"label\"].item()\n\n        showing_predictions = \"prediction\" in sample\n        if showing_predictions:\n            prediction = sample[\"prediction\"].item()\n\n        fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n\n        ax.imshow(image.permute(1, 2, 0)[:, :, band_idx])\n        ax.axis(\"off\")\n\n        if show_titles:\n            title = f\"Label: {label:.3f}\"\n            if showing_predictions:\n                title += f\"\\nPrediction: {prediction:.3f}\"\n            ax.set_title(title)\n\n        if suptitle is not None:\n            plt.suptitle(suptitle)\n\n        return fig\n", "path": "torchgeo/datasets/sustainbench_crop_yield.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"SustainBench Crop Yield dataset.\"\"\"\n\nimport os\nfrom typing import Any, Callable, Optional\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom matplotlib.figure import Figure\nfrom torch import Tensor\n\nfrom .geo import NonGeoDataset\nfrom .utils import DatasetNotFoundError, download_url, extract_archive\n\n\nclass SustainBenchCropYield(NonGeoDataset):\n    \"\"\"SustainBench Crop Yield Dataset.\n\n    This dataset contains MODIS band histograms and soybean yield\n    estimates for selected counties in the USA, Argentina and Brazil.\n    The dataset is part of the\n    `SustainBench <https://sustainlab-group.github.io/sustainbench/docs/datasets/sdg2/crop_yield.html>`_\n    datasets for tackling the UN Sustainable Development Goals (SDGs).\n\n    Dataset Format:\n\n    * .npz files of stacked samples\n\n    Dataset Features:\n\n    * input histogram of 7 surface reflectance and 2 surface temperature\n      bands from MODIS pixel values in 32 ranges across 32 timesteps\n      resulting in 32x32x9 input images\n    * regression target value of soybean yield in metric tonnes per\n      harvested hectare\n\n    If you use this dataset in your research, please cite:\n\n    * https://doi.org/10.1145/3209811.3212707\n    * https://doi.org/10.1609/aaai.v31i1.11172\n\n    .. versionadded:: 0.5\n    \"\"\"  # noqa: E501\n\n    valid_countries = [\"usa\", \"brazil\", \"argentina\"]\n\n    md5 = \"c2794e59512c897d9bea77b112848122\"\n\n    url = \"https://drive.google.com/file/d/1odwkI1hiE5rMZ4VfM0hOXzlFR4NbhrfU/view?usp=share_link\"  # noqa: E501\n\n    dir = \"soybeans\"\n\n    valid_splits = [\"train\", \"dev\", \"test\"]\n\n    def __init__(\n        self,\n        root: str = \"data\",\n        split: str = \"train\",\n        countries: list[str] = [\"usa\"],\n        transforms: Optional[Callable[[dict[str, Any]], dict[str, Any]]] = None,\n        download: bool = False,\n        checksum: bool = False,\n    ) -> None:\n        \"\"\"Initialize a new Dataset instance.\n\n        Args:\n            root: root directory where dataset can be found\n            split: one of \"train\", \"dev\", or \"test\"\n            countries: which countries to include in the dataset\n            transforms: a function/transform that takes an input sample\n                and returns a transformed version\n            download: if True, download dataset and store it in the root directory\n            checksum: if True, check the MD5 after downloading files (may be slow)\n\n        Raises:\n            AssertionError: if ``countries`` contains invalid countries or if ``split``\n                is invalid\n            DatasetNotFoundError: If dataset is not found and *download* is False.\n        \"\"\"\n        assert set(countries).issubset(\n            self.valid_countries\n        ), f\"Please choose a subset of these valid countried: {self.valid_countries}.\"\n        self.countries = countries\n\n        assert (\n            split in self.valid_splits\n        ), f\"Pleas choose one of these valid data splits {self.valid_splits}.\"\n        self.split = split\n\n        self.root = root\n        self.transforms = transforms\n        self.download = download\n        self.checksum = checksum\n\n        self._verify()\n        self.collection = self.retrieve_collection()\n\n    def __len__(self) -> int:\n        \"\"\"Return the number of data points in the dataset.\n\n        Returns:\n            length of the dataset\n        \"\"\"\n        return len(self.collection)\n\n    def __getitem__(self, index: int) -> dict[str, Tensor]:\n        \"\"\"Return an index within the dataset.\n\n        Args:\n            index: index to return\n\n        Returns:\n            data and label at that index\n        \"\"\"\n        input_file_path, sample_idx = self.collection[index]\n\n        sample: dict[str, Tensor] = {\n            \"image\": self._load_image(input_file_path, sample_idx)\n        }\n        sample.update(self._load_features(input_file_path, sample_idx))\n\n        if self.transforms is not None:\n            sample = self.transforms(sample)\n\n        return sample\n\n    def _load_image(self, path: str, sample_idx: int) -> Tensor:\n        \"\"\"Load input image.\n\n        Args:\n            path: path to input npz collection\n            sample_idx: what sample to index from the npz collection\n\n        Returns:\n            input image as tensor\n        \"\"\"\n        arr = np.load(path)[\"data\"][sample_idx]\n        # return [channel, height, width]\n        return torch.from_numpy(arr).permute(2, 0, 1).to(torch.float32)\n\n    def _load_features(self, path: str, sample_idx: int) -> dict[str, Tensor]:\n        \"\"\"Load features value.\n\n        Args:\n            path: path to image npz collection\n            sample_idx: what sample to index from the npz collection\n\n        Returns:\n            target regression value\n        \"\"\"\n        target_file_path = path.replace(\"_hists\", \"_yields\")\n        target = np.load(target_file_path)[\"data\"][sample_idx]\n\n        years_file_path = path.replace(\"_hists\", \"_years\")\n        year = int(np.load(years_file_path)[\"data\"][sample_idx])\n\n        ndvi_file_path = path.replace(\"_hists\", \"_ndvi\")\n        ndvi = np.load(ndvi_file_path)[\"data\"][sample_idx]\n\n        features = {\n            \"label\": torch.tensor(target).to(torch.float32),\n            \"year\": torch.tensor(year),\n            \"ndvi\": torch.from_numpy(ndvi).to(dtype=torch.float32),\n        }\n        return features\n\n    def retrieve_collection(self) -> list[tuple[str, int]]:\n        \"\"\"Retrieve the collection.\n\n        Returns:\n            path and index to dataset samples\n        \"\"\"\n        collection = []\n        for country in self.countries:\n            file_path = os.path.join(\n                self.root, self.dir, country, f\"{self.split}_hists.npz\"\n            )\n            npz_file = np.load(file_path)\n            num_data_points = npz_file[\"data\"].shape[0]\n            for idx in range(num_data_points):\n                collection.append((file_path, idx))\n\n        return collection\n\n    def _verify(self) -> None:\n        \"\"\"Verify the integrity of the dataset.\"\"\"\n        # Check if the extracted files already exist\n        pathname = os.path.join(self.root, self.dir)\n        if os.path.exists(pathname):\n            return\n\n        # Check if the zip files have already been downloaded\n        pathname = os.path.join(self.root, self.dir) + \".zip\"\n        if os.path.exists(pathname):\n            self._extract()\n            return\n\n        # Check if the user requested to download the dataset\n        if not self.download:\n            raise DatasetNotFoundError(self)\n\n        # Download the dataset\n        self._download()\n        self._extract()\n\n    def _download(self) -> None:\n        \"\"\"Download the dataset and extract it.\"\"\"\n        download_url(\n            self.url,\n            self.root,\n            filename=self.dir + \".zip\",\n            md5=self.md5 if self.checksum else None,\n        )\n        self._extract()\n\n    def _extract(self) -> None:\n        \"\"\"Extract the dataset.\"\"\"\n        zipfile_path = os.path.join(self.root, self.dir) + \".zip\"\n        extract_archive(zipfile_path, self.root)\n\n    def plot(\n        self,\n        sample: dict[str, Any],\n        band_idx: int = 0,\n        show_titles: bool = True,\n        suptitle: Optional[str] = None,\n    ) -> Figure:\n        \"\"\"Plot a sample from the dataset.\n\n        Args:\n            sample: a sample return by :meth:`__getitem__`\n            band_idx: which of the nine histograms to index\n            show_titles: flag indicating whether to show titles above each panel\n            suptitle: optional suptitle to use for figure\n\n        Returns:\n            a matplotlib Figure with the rendered sample\n\n        \"\"\"\n        image, label = sample[\"image\"], sample[\"label\"].item()\n\n        showing_predictions = \"prediction\" in sample\n        if showing_predictions:\n            prediction = sample[\"prediction\"].item()\n\n        fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n\n        ax.imshow(image.permute(1, 2, 0)[:, :, band_idx])\n        ax.axis(\"off\")\n\n        if show_titles:\n            title = f\"Label: {label:.3f}\"\n            if showing_predictions:\n                title += f\"\\nPrediction: {prediction:.3f}\"\n            ax.set_title(title)\n\n        if suptitle is not None:\n            plt.suptitle(suptitle)\n\n        return fig\n", "path": "torchgeo/datasets/sustainbench_crop_yield.py"}]}