Merge pull request #170 from jdb78/feature/improved-encoder-keyerror

Better KeyError for NaNLabelEncoder

Author: jdb78

pytorch_forecasting/data/encoders.py

@@ -102,7 +102,12 @@ def transform(self, y: Iterable) -> Union[torch.Tensor, np.ndarray]:
             encoded = [self.classes_.get(v, 0) for v in y]
 
         else:
-            encoded = [self.classes_[v] for v in y]
+            try:
+                encoded = [self.classes_[v] for v in y]
+            except KeyError as e:
+                raise KeyError(
+                    f"Unknown category '{e.args[0]}' encountered. Set `add_nan=True` to allow unknown categories"
+                )
 
         if isinstance(y, torch.Tensor):
             encoded = torch.tensor(encoded, dtype=torch.long, device=y.device)