Fix tests - ensure mase is calculated in correct places

Author: Jan Beitner

pytorch_forecasting/data/encoders.py

@@ -94,7 +94,10 @@ def transform(self, y: Iterable) -> Union[torch.Tensor, np.ndarray]:
             if self.warn:
                 cond = ~np.isin(y, self.classes_)
                 if cond.any():
-                    warnings.warn(f"Found {y[cond].nunique()} unknown classes which were set to NaN", UserWarning)
+                    warnings.warn(
+                        f"Found {np.unique(np.asarray(y)[cond]).size} unknown classes which were set to NaN",
+                        UserWarning,
+                    )
 
             encoded = [self.classes_.get(v, 0) for v in y]
 

pytorch_forecasting/models/base_model.py

@@ -183,11 +183,11 @@ def step(self, x: Dict[str, torch.Tensor], y: torch.Tensor, batch_idx: int, labe
             # for smoothness of loss function
             monotinicity_loss = 10 * torch.pow(monotinicity_loss, 2)
             if isinstance(self.loss, MASE):
-                self.loss(prediction, y)
-            else:
                 loss = self.loss(
                     prediction, y, encoder_target=x["encoder_target"], encoder_lengths=x["encoder_lengths"]
                 )
+            else:
+                loss = self.loss(prediction, y)
 
             loss = loss * (1 + monotinicity_loss)
         else:

pytorch_forecasting/models/nbeats/__init__.py

@@ -218,7 +218,10 @@ def step(self, x, y, batch_idx, label) -> Dict[str, torch.Tensor]:
             )
             backcast_weight = backcast_weight / (backcast_weight + 1)  # normalize
             forecast_weight = 1 - backcast_weight
-            backcast_loss = self.loss(backcast, x["encoder_target"]) * backcast_weight
+            if isinstance(self.loss, MASE):
+                backcast_loss = self.loss(backcast, x["encoder_target"], x["decoder_target"]) * backcast_weight
+            else:
+                backcast_loss = self.loss(backcast, x["encoder_target"]) * backcast_weight
             if label == "train":
                 log["loss"] = log["loss"] * forecast_weight + backcast_loss
                 log["log"]["train_loss"] = log["log"]["train_loss"] * forecast_weight + backcast_loss

pytorch_forecasting/models/temporal_fusion_transformer/sub_modules.py

@@ -351,15 +351,15 @@ def forward(self, x: Dict[str, torch.Tensor], context: torch.Tensor = None):
                     variable_embedding = self.prescalers[name](variable_embedding)
                 weight_inputs.append(variable_embedding)
                 var_outputs.append(self.single_variable_grns[name](variable_embedding))
-            var_outputs = torch.stack(var_outputs, axis=-1)
+            var_outputs = torch.stack(var_outputs, dim=-1)
 
             # calculate variable weights
             flat_embedding = torch.cat(weight_inputs, dim=-1)
             sparse_weights = self.flattened_grn(flat_embedding, context)
             sparse_weights = self.softmax(sparse_weights).unsqueeze(-2)
 
             outputs = var_outputs * sparse_weights
-            outputs = outputs.sum(axis=-1)
+            outputs = outputs.sum(dim=-1)
         else:  # for one input, do not perform variable selection but just encoding
             name = next(iter(self.single_variable_grns.keys()))
             variable_embedding = x[name]