diff --git a/etsformer/lightning_module.py b/etsformer/lightning_module.py
index 184d406..43ac3c4 100644
--- a/etsformer/lightning_module.py
+++ b/etsformer/lightning_module.py
@@ -75,6 +75,6 @@ class ETSformerLightningModule(pl.LightningModule):
         if len(self.model.target_shape) == 0:
             loss_weights = future_observed_values
         else:
-            loss_weights = future_observed_values.min(dim=-1, keepdim=False)
+            loss_weights, _ = future_observed_values.min(dim=-1, keepdim=False)
 
         return weighted_average(loss_values, weights=loss_weights)
diff --git a/etsformer/module.py b/etsformer/module.py
index 6901b7c..b988c1a 100644
--- a/etsformer/module.py
+++ b/etsformer/module.py
@@ -85,7 +85,7 @@ class ETSformerModel(nn.Module):
             sum(self.embedding_dimension)
             + self.num_feat_dynamic_real
             + self.num_feat_static_real
-            + 1  # the log(scale)
+            + self.input_size  # the log(scale)
         )
 
     @property
@@ -194,8 +194,9 @@ class ETSformerModel(nn.Module):
 
         # embeddings
         embedded_cat = self.embedder(feat_static_cat)
+        log_scale = scale.log() if self.input_size == 1 else scale.squeeze(1).log()
         static_feat = torch.cat(
-            (embedded_cat, feat_static_real, scale.log()),
+            (embedded_cat, feat_static_real, log_scale),
             dim=1,
         )
         expanded_static_feat = static_feat.unsqueeze(1).expand(