# Buggy code
def save(self, dir_path: str, overwrite: bool = False):
    """
    Save the state of the model.

    Neither the trainer optimizer state nor the trainer history are saved.
    Model files are not expected to be reproducibly saved and loaded across versions
    until we reach version 1.0.

    Parameters
    ----------
    dir_path
        Path to a directory.
    overwrite
        Overwrite existing data or not. If `False` and directory
        already exists at `dir_path`, error will be raised.
    """
    # get all the user attributes
    user_attributes = self._get_user_attributes()
    # only save the public attributes with _ at the very end
    user_attributes = {a[0]: a[1] for a in user_attributes if a[0][-1] == "_"}
    # save the model state dict and the trainer state dict only
    if not os.path.exists(dir_path) or overwrite:
        os.makedirs(dir_path, exist_ok=overwrite)
    else:
        raise ValueError(
            "{} already exists. Please provide an unexisting directory for saving.".format(
                dir_path
            )
        )
    torch.save(self.model.state_dict(), os.path.join(dir_path, "model_params.pt"))
    with open(os.path.join(dir_path, "attr.pkl"), "wb") as f:
        pickle.dump(user_attributes, f)


# 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 save(
    self,
    dir_path: str,
    overwrite: bool = False,
    save_anndata: bool = False,
    **anndata_write_kwargs,
):
    """
    Save the state of the model.

    Neither the trainer optimizer state nor the trainer history are saved.
    Model files are not expected to be reproducibly saved and loaded across versions
    until we reach version 1.0.

    Parameters
    ----------
    dir_path
        Path to a directory.
    overwrite
        Overwrite existing data or not. If `False` and directory
        already exists at `dir_path`, error will be raised.
    save_anndata
        If True, also saves the anndata
    anndata_write_kwargs
        Kwargs for anndata write function
    """
    # get all the user attributes
    user_attributes = self._get_user_attributes()
    # only save the public attributes with _ at the very end
    user_attributes = {a[0]: a[1] for a in user_attributes if a[0][-1] == "_"}
    # save the model state dict and the trainer state dict only
    if not os.path.exists(dir_path) or overwrite:
        os.makedirs(dir_path, exist_ok=overwrite)
    else:
        raise ValueError(
            "{} already exists. Please provide an unexisting directory for saving.".format(
                dir_path
            )
        )

    if save_anndata:
        self.adata.write(os.path.join(dir_path, "adata.h5ad"), **anndata_write_kwargs)

    model_save_path = os.path.join(dir_path, "model_params.pt")
    attr_save_path = os.path.join(dir_path, "attr.pkl")
    varnames_save_path = os.path.join(dir_path, "var_names.csv")

    var_names = self.adata.var_names.astype(str)
    var_names = var_names.to_numpy()
    np.savetxt(varnames_save_path, var_names, fmt="%s")

    torch.save(self.model.state_dict(), model_save_path)
    with open(attr_save_path, "wb") as f:
        pickle.dump(user_attributes, f)