diff --git a/velovi/_model.py b/velovi/_model.py
index e75b98b..e5f1115 100644
--- a/velovi/_model.py
+++ b/velovi/_model.py
@@ -49,6 +49,7 @@ class VELOVI(VAEMixin, UnsupervisedTrainingMixin, BaseModelClass):
         Use a linear decoder from latent space to time.
     **model_kwargs
         Keyword args for :class:`~velovi.VELOVAE`
+
     """
 
     def __init__(
@@ -163,6 +164,7 @@ def train(
             `train()` will overwrite values present in `plan_kwargs`, when appropriate.
         **trainer_kwargs
             Other keyword args for :class:`~scvi.train.Trainer`.
+
         """
         user_plan_kwargs = plan_kwargs.copy() if isinstance(plan_kwargs, dict) else {}
         plan_kwargs = {"lr": lr, "weight_decay": weight_decay, "optimizer": "AdamW"}
@@ -236,6 +238,7 @@ def get_state_assignment(
         -------
         If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
         Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
+
         """
         adata = self._validate_anndata(adata)
         scdl = self._make_data_loader(
@@ -340,6 +343,7 @@ def get_latent_time(
         -------
         If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
         Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
+
         """
         adata = self._validate_anndata(adata)
         if indices is None:
@@ -482,6 +486,7 @@ def get_velocity(
         -------
         If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
         Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
+
         """
         adata = self._validate_anndata(adata)
         if indices is None:
@@ -656,6 +661,7 @@ def get_expression_fit(
         -------
         If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
         Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
+
         """
         adata = self._validate_anndata(adata)
 
@@ -811,6 +817,7 @@ def get_gene_likelihood(
         -------
         If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
         Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
+
         """
         adata = self._validate_anndata(adata)
         scdl = self._make_data_loader(
@@ -917,6 +924,7 @@ def setup_anndata(
         Returns
         -------
         %(returns)s
+
         """
         setup_method_args = cls._get_setup_method_args(**locals())
         anndata_fields = [
@@ -967,6 +975,7 @@ def get_permutation_scores(
         -------
         Tuple of DataFrame and AnnData. DataFrame is genes by cell types with score per cell type.
         AnnData is the permutated version of the original AnnData.
+
         """
         adata = self._validate_anndata(adata)
         adata_manager = self.get_anndata_manager(adata)
@@ -1090,6 +1099,7 @@ def _directional_statistics_per_cell(
     ----------
     tensor
         Shape of samples by genes for a given cell.
+
     """
     n_samples = tensor.shape[0]
     # over samples axis
diff --git a/velovi/_module.py b/velovi/_module.py
index 2515b16..9a7ec9a 100644
--- a/velovi/_module.py
+++ b/velovi/_module.py
@@ -1,4 +1,5 @@
 """Main module."""
+
 from typing import Callable, Iterable, Literal, Optional
 
 import numpy as np
@@ -44,6 +45,7 @@ class DecoderVELOVI(nn.Module):
         Whether to use layer norm in layers
     linear_decoder
         Whether to use linear decoder for time
+
     """
 
     def __init__(
@@ -183,6 +185,7 @@ class VELOVAE(BaseModuleClass):
     var_activation
         Callable used to ensure positivity of the variational distributions' variance.
         When `None`, defaults to `torch.exp`.
+
     """
 
     def __init__(
diff --git a/velovi/_utils.py b/velovi/_utils.py
index bfffbeb..d16cc48 100644
--- a/velovi/_utils.py
+++ b/velovi/_utils.py
@@ -56,6 +56,7 @@ def preprocess_data(
     Returns
     -------
     Preprocessed adata.
+
     """
     if min_max_scale:
         scaler = MinMaxScaler()