# Buggy code
def save(self, labels: SemanticSegmentationLabels) -> None:
    """Save labels to disk.

    More info on rasterio IO:
    - https://github.com/mapbox/rasterio/blob/master/docs/quickstart.rst
    - https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html

    Args:
        labels - (SemanticSegmentationLabels) labels to be saved
    """
    local_root = get_local_path(self.root_uri, self.tmp_dir)
    make_dir(local_root)

    out_smooth_profile = {
        "driver": "GTiff",
        "height": self.extent.ymax,
        "width": self.extent.xmax,
        "transform": self.crs_transformer.get_affine_transform(),
        "crs": self.crs_transformer.get_image_crs(),
        "blockxsize": self.rasterio_block_size,
        "blockysize": self.rasterio_block_size,
    }

    # if old scores exist, combine them with the new ones
    if self.score_raster_source:
        log.info("Old scores found. Merging with current scores.")
        old_labels = self.get_scores()
        labels += old_labels

    self.write_discrete_raster_output(
        out_smooth_profile, get_local_path(self.label_uri, self.tmp_dir), labels
    )

    if self.smooth_output:
        self.write_smooth_raster_output(
            out_smooth_profile,
            get_local_path(self.score_uri, self.tmp_dir),
            get_local_path(self.hits_uri, self.tmp_dir),
            labels,
            chip_sz=self.rasterio_block_size,
        )

    if self.vector_outputs:
        self.write_vector_outputs(labels)

    sync_to_dir(local_root, self.root_uri)


# Traceback (KeyError)
Running predict command...
2021-01-18 09:29:36:rastervision.pytorch_learner.learner: INFO - Loading model weights from: /opt/data/tmp/tmpno2kuo4_/model-bundle/model.pth
2021-01-18 09:29:38:rastervision.core.rv_pipeline.rv_pipeline: INFO - Making predictions for scene
................................................................................................................................................................................................................................................................................................................................................................................................................
2021-01-18 09:29:48:rastervision.core.data.label_store.semantic_segmentation_label_store: INFO - Writing labels to disk.
[#-----------------------------------]    5%  0d 00:00:48
Traceback (most recent call last):
File "/opt/conda/lib/python3.6/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/opt/conda/lib/python3.6/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 248, in <module>
main()
File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 722, in __call__
return self.main(*args, **kwargs)
File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 697, in main
rv = self.invoke(ctx)
File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 1066, in invoke
return _process_result(sub_ctx.command.invoke(sub_ctx))
File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 895, in invoke
return ctx.invoke(self.callback, **ctx.params)
File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 535, in invoke
return callback(*args, **kwargs)
File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 240, in run_command
runner=runner)
File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 217, in _run_command
command_fn()
File "/opt/src/rastervision_core/rastervision/core/rv_pipeline/rv_pipeline.py", line 165, in predict
for s in dataset.validation_scenes
File "/opt/src/rastervision_core/rastervision/core/rv_pipeline/rv_pipeline.py", line 161, in _predict
label_store.save(labels)
File "/opt/src/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py", line 194, in save
labels)
File "/opt/src/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py", line 256, in write_discrete_raster_output
label_arr = labels.get_label_arr(window)
File "/opt/src/rastervision_core/rastervision/core/data/label/semantic_segmentation_labels.py", line 170, in get_label_arr
return self.window_to_label_arr[window]
KeyError: Box(0, 4864, 256, 5000)

# Fixed code
def save(self, labels: SemanticSegmentationLabels) -> None:
    """Save labels to disk.

    More info on rasterio IO:
    - https://github.com/mapbox/rasterio/blob/master/docs/quickstart.rst
    - https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html

    Args:
        labels - (SemanticSegmentationLabels) labels to be saved
    """
    local_root = get_local_path(self.root_uri, self.tmp_dir)
    make_dir(local_root)

    out_profile = {
        "driver": "GTiff",
        "height": self.extent.ymax,
        "width": self.extent.xmax,
        "transform": self.crs_transformer.get_affine_transform(),
        "crs": self.crs_transformer.get_image_crs(),
        "blockxsize": self.rasterio_block_size,
        "blockysize": self.rasterio_block_size,
    }

    # if old scores exist, combine them with the new ones
    if self.score_raster_source:
        log.info("Old scores found. Merging with current scores.")
        old_labels = self.get_scores()
        labels += old_labels

    self.write_discrete_raster_output(
        out_profile, get_local_path(self.label_uri, self.tmp_dir), labels
    )

    if self.smooth_output:
        self.write_smooth_raster_output(
            out_profile,
            get_local_path(self.score_uri, self.tmp_dir),
            get_local_path(self.hits_uri, self.tmp_dir),
            labels,
            chip_sz=self.rasterio_block_size,
        )

    if self.vector_outputs:
        self.write_vector_outputs(labels)

    sync_to_dir(local_root, self.root_uri)