# Buggy code
def write_discrete_raster_output(
    self, out_smooth_profile: dict, path: str, labels: SemanticSegmentationLabels
) -> None:
    num_bands = 1 if self.class_transformer is None else 3
    out_smooth_profile.update({"count": num_bands, "dtype": np.uint8})

    windows = labels.get_windows()

    log.info("Writing labels to disk.")
    with rio.open(path, "w", **out_smooth_profile) as dataset:
        with click.progressbar(windows) as bar:
            for window in bar:
                window = window.intersection(self.extent)
                label_arr = labels.get_label_arr(window)
                window = window.rasterio_format()
                if self.class_transformer is None:
                    dataset.write_band(1, label_arr, window=window)
                else:
                    rgb_labels = self.class_transformer.class_to_rgb(label_arr)
                    rgb_labels = rgb_labels.transpose(2, 0, 1)
                    for i, band in enumerate(rgb_labels, start=1):
                        dataset.write_band(i, band, window=window)


# 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 write_discrete_raster_output(
    self, out_profile: dict, path: str, labels: SemanticSegmentationLabels
) -> None:
    num_bands = 1 if self.class_transformer is None else 3
    out_profile.update({"count": num_bands, "dtype": np.uint8})

    windows = labels.get_windows()

    log.info("Writing labels to disk.")
    with rio.open(path, "w", **out_profile) as dataset:
        with click.progressbar(windows) as bar:
            for window in bar:
                label_arr = labels.get_label_arr(window)
                window, label_arr = self._clip_to_extent(self.extent, window, label_arr)
                if self.class_transformer is not None:
                    label_arr = self.class_transformer.class_to_rgb(label_arr)
                    label_arr = label_arr.transpose(2, 0, 1)
                self._write_array(dataset, window, label_arr)