diff --git a/pytorch_forecasting/layers/_decomposition/_series_decomp.py b/pytorch_forecasting/layers/_decomposition/_series_decomp.py
index 30c6b8b38..15fb006b3 100644
--- a/pytorch_forecasting/layers/_decomposition/_series_decomp.py
+++ b/pytorch_forecasting/layers/_decomposition/_series_decomp.py
@@ -1,6 +1,4 @@
-"""
-Series Decomposition Block for time series forecasting models.
-"""
+"""Series Decomposition Block for time series forecasting models."""
 
 import torch
 import torch.nn as nn
@@ -10,15 +8,15 @@
 
 
 class SeriesDecomposition(nn.Module):
-    """
-    Series decomposition block from Autoformer.
+    """Series decomposition block from Autoformer.
 
     Decomposes time series into trend and seasonal components using
     moving average filtering.
 
-    Args:
-        kernel_size (int):
-            Size of the moving average kernel for trend extraction.
+    Parameters
+    ----------
+    kernel_size : int
+        Size of the moving average kernel for trend extraction.
     """
 
     def __init__(self, kernel_size):
@@ -26,17 +24,20 @@ def __init__(self, kernel_size):
         self.moving_avg = MovingAvg(kernel_size, stride=1)
 
     def forward(self, x):
-        """
-        Forward pass for series decomposition.
-
-        Args:
-            x (torch.Tensor):
-                Input time series tensor of shape (batch_size, seq_len, features).
-
-        Returns:
-            tuple:
-                - trend (torch.Tensor): Trend component of the time series.
-                - seasonal (torch.Tensor): Seasonal component of the time series.
+        """Forward pass for series decomposition.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input time series tensor of shape (batch_size, seq_len, features).
+
+        Returns
+        -------
+        tuple
+            seasonal : torch.Tensor
+                Seasonal component of the time series.
+            trend : torch.Tensor
+                Trend component of the time series.
         """
         trend = self.moving_avg(x)
         seasonal = x - trend
diff --git a/pytorch_forecasting/layers/_output/_flatten_head.py b/pytorch_forecasting/layers/_output/_flatten_head.py
index ba2aa0eae..8e1e9e8a7 100644
--- a/pytorch_forecasting/layers/_output/_flatten_head.py
+++ b/pytorch_forecasting/layers/_output/_flatten_head.py
@@ -1,6 +1,4 @@
-"""
-Implementation of output layers from `nn.Module` for TimeXer model.
-"""
+"""Implementation of output layers from `nn.Module` for TimeXer model."""
 
 import math
 from math import sqrt
@@ -12,14 +10,20 @@
 
 
 class FlattenHead(nn.Module):
-    """
-    Flatten head for the output of the model.
-    Args:
-        n_vars (int): Number of input features.
-        nf (int): Number of features in the last layer.
-        target_window (int): Target window size.
-        head_dropout (float): Dropout rate for the head. Defaults to 0.
-        n_quantiles (int, optional): Number of quantiles. Defaults to None."""
+    """Flatten head for the output of the model.
+
+    Parameters
+    ----------
+    n_vars : int
+        Number of input features.
+    nf : int
+        Number of features in the last layer.
+        target_window : int
+        Target window size.
+    head_dropout : float, optional
+        Dropout rate for the head. Defaults to 0.
+    n_quantiles : int, optional
+        Number of quantiles. Defaults to None."""
 
     def __init__(self, n_vars, nf, target_window, head_dropout=0, n_quantiles=None):
         super().__init__()