# Buggy code
def train(
    self,
    n_epochs: Optional[int] = None,
    train_size: float = 0.9,
    test_size: Optional[float] = None,
    lr: float = 1e-3,
    n_epochs_kl_warmup: int = 400,
    n_iter_kl_warmup: Optional[int] = None,
    frequency: Optional[int] = None,
    train_fun_kwargs: dict = {},
    **kwargs,
):
    """
    Trains the model using amortized variational inference.

    Parameters
    ----------
    n_epochs
        Number of passes through the dataset.
    train_size
        Size of training set in the range [0.0, 1.0].
    test_size
        Size of the test set. If `None`, defaults to 1 - `train_size`. If
        `train_size + test_size < 1`, the remaining cells belong to a validation set.
    lr
        Learning rate for optimization.
    n_epochs_kl_warmup
        Number of passes through dataset for scaling term on KL divergence to go from 0 to 1.
    n_iter_kl_warmup
        Number of minibatches for scaling term on KL divergence to go from 0 to 1.
        To use, set to not `None` and set `n_epochs_kl_warmup` to `None`.
    frequency
        Frequency with which metrics are computed on the data for train/test/val sets.
    train_fun_kwargs
        Keyword args for the train method of :class:`~scvi.core.trainers.UnsupervisedTrainer`.
    **kwargs
        Other keyword args for :class:`~scvi.core.trainers.UnsupervisedTrainer`.
    """
    train_fun_kwargs = dict(train_fun_kwargs)
    if self.is_trained_ is False:
        self.trainer = UnsupervisedTrainer(
            self.model,
            self.adata,
            train_size=train_size,
            test_size=test_size,
            n_iter_kl_warmup=n_iter_kl_warmup,
            n_epochs_kl_warmup=n_epochs_kl_warmup,
            frequency=frequency,
            use_cuda=self.use_cuda,
            **kwargs,
        )
        self.train_indices_ = self.trainer.train_set.indices
        self.test_indices_ = self.trainer.test_set.indices
        self.validation_indices_ = self.trainer.validation_set.indices
    # for autotune
    if "n_epochs" not in train_fun_kwargs:
        if n_epochs is None:
            n_cells = self.adata.n_obs
            n_epochs = np.min([round((20000 / n_cells) * 400), 400])
        train_fun_kwargs["n_epochs"] = n_epochs
    if "lr" not in train_fun_kwargs:
        train_fun_kwargs["lr"] = lr

    logger.info("Training for {} epochs".format(n_epochs))
    self.trainer.train(**train_fun_kwargs)
    self.is_trained_ = True


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-125-eb1e0fcfe1d1> in <module>
1 de = model.differential_expression(
----> 2     groupby="leiden",
3 )
4 de.head()

~/anaconda3/lib/python3.7/site-packages/scvi/core/models/rnamixin.py in differential_expression(self, adata, groupby, group1, group2, idx1, idx2, mode, delta, batch_size, all_stats, batch_correction, batchid1, batchid2, **kwargs)
179         adata = self._validate_anndata(adata)
180
--> 181         col_names = _get_var_names_from_setup_anndata(adata)
182         model_fn = partial(
183             self.get_normalized_expression,

~/anaconda3/lib/python3.7/site-packages/scvi/model/_utils.py in _get_var_names_from_setup_anndata(adata)
119         adata.var_names
120         if adata.uns["_scvi"]["use_raw"] is False
--> 121         else adata.raw.var_names
122     )
123

AttributeError: 'NoneType' object has no attribute 'var_names'

# Fixed code
def train(
    self,
    n_epochs: Optional[int] = None,
    train_size: float = 0.9,
    test_size: Optional[float] = None,
    lr: float = 1e-3,
    n_epochs_kl_warmup: int = 400,
    n_iter_kl_warmup: Optional[int] = None,
    frequency: Optional[int] = None,
    train_fun_kwargs: dict = {},
    **kwargs,
):
    """
    Trains the model using amortized variational inference.

    Parameters
    ----------
    n_epochs
        Number of passes through the dataset.
    train_size
        Size of training set in the range [0.0, 1.0].
    test_size
        Size of the test set. If `None`, defaults to 1 - `train_size`. If
        `train_size + test_size < 1`, the remaining cells belong to a validation set.
    lr
        Learning rate for optimization.
    n_epochs_kl_warmup
        Number of passes through dataset for scaling term on KL divergence to go from 0 to 1.
    n_iter_kl_warmup
        Number of minibatches for scaling term on KL divergence to go from 0 to 1.
        To use, set to not `None` and set `n_epochs_kl_warmup` to `None`.
    frequency
        Frequency with which metrics are computed on the data for train/test/val sets.
    train_fun_kwargs
        Keyword args for the train method of :class:`~scvi.core.trainers.UnsupervisedTrainer`.
    **kwargs
        Other keyword args for :class:`~scvi.core.trainers.UnsupervisedTrainer`.
    """
    train_fun_kwargs = dict(train_fun_kwargs)
    if self.is_trained_ is False:
        self.trainer = UnsupervisedTrainer(
            self.model,
            self.adata,
            train_size=train_size,
            test_size=test_size,
            n_iter_kl_warmup=n_iter_kl_warmup,
            n_epochs_kl_warmup=n_epochs_kl_warmup,
            frequency=frequency,
            use_cuda=self.use_cuda,
            **kwargs,
        )
        self.train_indices_ = self.trainer.train_set.indices
        self.test_indices_ = self.trainer.test_set.indices
        self.validation_indices_ = self.trainer.validation_set.indices
        self.history_ = self.trainer.history
    # for autotune
    if "n_epochs" not in train_fun_kwargs:
        if n_epochs is None:
            n_cells = self.adata.n_obs
            n_epochs = np.min([round((20000 / n_cells) * 400), 400])
        train_fun_kwargs["n_epochs"] = n_epochs
    if "lr" not in train_fun_kwargs:
        train_fun_kwargs["lr"] = lr

    logger.info("Training for {} epochs".format(n_epochs))
    self.trainer.train(**train_fun_kwargs)
    self.is_trained_ = True