Fix label_override when forcing_loader not given labels (ai2cm#977)

During ace inference data loading, if forcing_loader does not have
labels set, label_override breaks due to the lack of a label_encoding in
the forcing dataset (see [beaker
log](https://beaker.org/orgs/ai2/workspaces/ace/work/01KKN10MF6Q1NJJFE10GT099X4/logs?jobId=01KKN10MS6DX9GC4VFKHQ32J0T)).
This PR resolves the issue by using label_override to define the
label_encoding when it is used. This avoids the awkward position of
needing to set both the label at the highest level of InferenceConfig
and in the forcing dataset.

What this PR doesn't address is the potential misconfiguration of labels
in `InferenceConfig` is set to ERA5, but the forcing dataset label was
mistakenly set to a different label. We should address this in a future
PR.

Changes:
- Create a `LabelEncoding` if the dataset label isn't set but the
`label_override` is not None
- Add test to `test_data_loader.py‎`

- [x] Tests added

Author: Arcomano1234

fme/ace/data_loading/inference.py

@@ -225,6 +225,10 @@ def __init__(
         """
         if label_encoding is None and config.available_labels is not None:
             label_encoding = LabelEncoding(labels=sorted(list(config.available_labels)))
+        elif label_encoding is None and label_override is not None:
+            # When labels are overridden (e.g. from config.labels), we still need
+            # an encoding to collate them even if the dataset has no available_labels.
+            label_encoding = LabelEncoding(labels=sorted(list(label_override)))
         self._label_encoding = label_encoding
         self._label_override = (
             set(label_override) if label_override is not None else None

fme/ace/data_loading/test_data_loader.py

@@ -962,6 +962,31 @@ def test_inference_persistence_names(tmp_path):
     assert not torch.all(first_item["bar"] == second_item["bar"])
 
 
+def test_inference_dataset_label_override_without_dataset_labels(tmp_path):
+    _create_dataset_on_disk(tmp_path, n_times=14)
+    config = InferenceDataLoaderConfig(
+        dataset=XarrayDataConfig(data_path=tmp_path),
+        start_indices=ExplicitIndices([0]),
+    )
+    window_requirements = DataRequirements(
+        names=["foo", "bar"],
+        n_timesteps=3,
+    )
+    dataset = InferenceDataset(
+        config,
+        total_forward_steps=3,
+        requirements=window_requirements,
+        label_override=["era5"],
+    )
+    batch = dataset[0]
+    # assert that the labels are not None and are the correct labels set during
+    # inference config initialization
+    assert batch.labels is not None
+    assert batch.labels.names == ["era5"]
+    assert batch.labels.tensor.shape[0] == 1
+    assert batch.labels.tensor.shape[1] == 1
+
+
 def test_zarr_engine_used_sequence():
     config = DataLoaderConfig(
         dataset=ConcatDatasetConfig(