diff --git a/doc/extras.inv b/doc/extras.inv
index 88973215f3227495564e345c34603dad42a9b532..d053cdcf2c67163e69d16b9ab3cff2f0aaaf4f37 100644
--- a/doc/extras.inv
+++ b/doc/extras.inv
@@ -2,5 +2,5 @@
 # Project: extras
 # Version: stable
 # The remainder of this file is compressed using zlib.
-xÚEËÁ
€ лSt
i¼²
*
-	PÒ~MØÞÞ߃è–îlYšƒ†f‡h5êÃWÙ¯i¡tóÌ}àÅNôäo°!¬%ò]B-4OÎŒ ã
\ No newline at end of file
+xÚ­A
+1E÷ž¢ [[ÜÎôR§±
´MH£2žÞ‡A\ˆ.ÃÏ{OIútņT­Š¯íLRšá¡+.ÌÎ$Uns<èôlI¢›	×ÔŸ2¸h“–l¶«Œ1iÅíBõ$`g§Ý/ëa¾Gôæ%<« ‰ÂXõŒîƒåŸšëÍ×UØë±)ðÏibÅ‚wÿèĘ/
\ No newline at end of file
diff --git a/doc/extras.txt b/doc/extras.txt
index e827f8fa4af7ea94634d044dec1d282029babcc4..77fd0ca6112dcd91c2c303eefe87b279f3f64035 100644
--- a/doc/extras.txt
+++ b/doc/extras.txt
@@ -3,3 +3,6 @@
 # Version: stable
 # The remainder of this file is compressed using zlib.
 torchvision.transforms py:module 1 https://pytorch.org/vision/stable/transforms.html -
+lightning.pytorch.core.module.LightningModule.forward py:method 1 api/lightning.pytorch.core.LightningModule.html#$ -
+lightning.pytorch.core.module.LightningModule.predict_step py:method 1 api/lightning.pytorch.core.LightningModule.html#$ -
+optimizer_step py:method 1 api/lightning.pytorch.core.LightningModule.html#$ -
diff --git a/src/ptbench/models/alexnet.py b/src/ptbench/models/alexnet.py
index 073013cd5a24ec370e8d594d0c4a7af050fef0b1..10ecfc7215ecd0903d848812669d8ae4debf9ff6 100644
--- a/src/ptbench/models/alexnet.py
+++ b/src/ptbench/models/alexnet.py
@@ -10,7 +10,7 @@ import torchvision.models as models
 from .normalizer import TorchVisionNormalizer
 
 
-class Alexnet(pl.core.LightningModule):
+class Alexnet(pl.LightningModule):
     """Alexnet module.
 
     Note: only usable with a normalized dataset
@@ -43,20 +43,6 @@ class Alexnet(pl.core.LightningModule):
         self.model_ft.classifier[6] = nn.Linear(512, 1)
 
     def forward(self, x):
-        """
-
-        Parameters
-        ----------
-
-        x : list
-            list of tensors.
-
-        Returns
-        -------
-
-        tensor : :py:class:`torch.Tensor`
-
-        """
         x = self.normalizer(x)
         x = self.model_ft(x)
 
diff --git a/src/ptbench/models/densenet.py b/src/ptbench/models/densenet.py
index 27c3393df25e8695f647bada0418e72d1c7f821f..77cbc0a8d3443a940489cc7f482655cd183c0cc5 100644
--- a/src/ptbench/models/densenet.py
+++ b/src/ptbench/models/densenet.py
@@ -43,22 +43,7 @@ class Densenet(pl.LightningModule):
         )
 
     def forward(self, x):
-        """
-
-        Parameters
-        ----------
-
-        x : list
-            list of tensors.
-
-        Returns
-        -------
-
-        tensor : :py:class:`torch.Tensor`
-
-        """
         x = self.normalizer(x)
-
         x = self.model_ft(x)
 
         return x
diff --git a/src/ptbench/models/densenet_rs.py b/src/ptbench/models/densenet_rs.py
index 16f4eefb2891e7a1185b3141d0f6369c0169642b..6e5a3df4f46db4d4b2fb70dfcdf55b1b2a7decfa 100644
--- a/src/ptbench/models/densenet_rs.py
+++ b/src/ptbench/models/densenet_rs.py
@@ -38,21 +38,6 @@ class DensenetRS(pl.LightningModule):
         self.model_ft.classifier = nn.Linear(num_ftrs, 14)
 
     def forward(self, x):
-        """
-
-        Parameters
-        ----------
-
-        x : list
-            list of tensors.
-
-        Returns
-        -------
-
-        tensor : :py:class:`torch.Tensor`
-
-        """
-
         x = self.normalizer(x)
         x = self.model_ft(x)
         return x
diff --git a/src/ptbench/models/logistic_regression.py b/src/ptbench/models/logistic_regression.py
index c6df54bcd8501de8dccb9d4cd475a0e4fa60d84f..485a396760facfc3bf84fbbaf1b66db40f327aa3 100644
--- a/src/ptbench/models/logistic_regression.py
+++ b/src/ptbench/models/logistic_regression.py
@@ -27,20 +27,6 @@ class LogisticRegression(pl.LightningModule):
         self.linear = nn.Linear(self.hparams.input_size, 1)
 
     def forward(self, x):
-        """
-
-        Parameters
-        ----------
-
-        x : list
-            list of tensors.
-
-        Returns
-        -------
-
-        tensor : :py:class:`torch.Tensor`
-
-        """
         output = self.linear(x)
 
         return output
diff --git a/src/ptbench/models/pasa.py b/src/ptbench/models/pasa.py
index 3d4a7641bd82c6bdf461b90bdfe735ef7a8b3332..d4e5b2a85203bc9aa6be2d1642e84b903497024e 100644
--- a/src/ptbench/models/pasa.py
+++ b/src/ptbench/models/pasa.py
@@ -93,20 +93,6 @@ class PASA(pl.LightningModule):
         self.dense = nn.Linear(80, 1)  # Fully connected layer
 
     def forward(self, x):
-        """
-
-        Parameters
-        ----------
-
-        x : list
-            list of tensors.
-
-        Returns
-        -------
-
-        tensor : :py:class:`torch.Tensor`
-
-        """
         x = self.normalizer(x)
 
         # First convolution block
diff --git a/src/ptbench/models/signs_to_tb.py b/src/ptbench/models/signs_to_tb.py
index 4733772780951dc7999946895ae24c5ee89a10c7..9267e7778dacc46d58972f016f79e25e0b47366a 100644
--- a/src/ptbench/models/signs_to_tb.py
+++ b/src/ptbench/models/signs_to_tb.py
@@ -31,20 +31,6 @@ class SignsToTB(pl.LightningModule):
         self.fc2 = torch.nn.Linear(self.hparams.hidden_size, 1)
 
     def forward(self, x):
-        """
-
-        Parameters
-        ----------
-
-        x : list
-            list of tensors.
-
-        Returns
-        -------
-
-        tensor : :py:class:`torch.Tensor`
-
-        """
         hidden = self.fc1(x)
         relu = self.relu(hidden)