# Buggy code
def _set_data_in_registry(adata, data, key):
    """
    Sets the data associated with key in adata.uns['_scvi']['data_registry'].keys() to data.

    Note: This is a dangerous method and will change the underlying data of the user's anndata
    Currently used to make the user's anndata C_CONTIGUOUS and csr if it is dense numpy
    or sparse respectively.

    Parameters
    ----------
    adata
        anndata object to change data of
    data
        data to change to
    key
        key in adata.uns['_scvi]['data_registry'].keys() associated with the data
    """
    use_raw = adata.uns["_scvi"]["use_raw"]
    data_loc = adata.uns["_scvi"]["data_registry"][key]
    attr_name, attr_key = data_loc["attr_name"], data_loc["attr_key"]

    if use_raw is True and attr_name in ["X", "var"]:
        tmp_adata = adata.raw.to_adata()
    else:
        tmp_adata = adata
    if attr_key == "None":
        setattr(tmp_adata, attr_name, data)

    elif attr_key != "None":
        attribute = getattr(tmp_adata, attr_name)
        if isinstance(attribute, pd.DataFrame):
            attribute.loc[:, attr_key] = data
        else:
            attribute[attr_key] = data
        setattr(tmp_adata, attr_name, attribute)

    if use_raw is True and attr_name in ["X", "var"]:
        adata.raw = tmp_adata


# 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 _set_data_in_registry(adata, data, key):
    """
    Sets the data associated with key in adata.uns['_scvi']['data_registry'].keys() to data.

    Note: This is a dangerous method and will change the underlying data of the user's anndata
    Currently used to make the user's anndata C_CONTIGUOUS and csr if it is dense numpy
    or sparse respectively.

    Parameters
    ----------
    adata
        anndata object to change data of
    data
        data to change to
    key
        key in adata.uns['_scvi]['data_registry'].keys() associated with the data
    """
    data_loc = adata.uns["_scvi"]["data_registry"][key]
    attr_name, attr_key = data_loc["attr_name"], data_loc["attr_key"]

    if attr_key == "None":
        setattr(adata, attr_name, data)

    elif attr_key != "None":
        attribute = getattr(adata, attr_name)
        if isinstance(attribute, pd.DataFrame):
            attribute.loc[:, attr_key] = data
        else:
            attribute[attr_key] = data
        setattr(adata, attr_name, attribute)