# Buggy code
def __init__(
    self,
    filename_or_anndata,
    save_path="data/",
    url=None,
    new_n_genes=False,
    subset_genes=None,
):
    """ """

    if type(filename_or_anndata) == str:
        self.download_name = filename_or_anndata
        self.save_path = save_path
        self.url = url

        data, gene_names, batch_indices, cell_types, labels = (
            self.download_and_preprocess()
        )

    elif isinstance(filename_or_anndata, anndata.AnnData):
        ad = filename_or_anndata
        data, gene_names, batch_indices, cell_types, labels = (
            self.extract_data_from_anndata(ad)
        )

    else:
        raise Exception(
            "Please provide a filename of an AnnData file or an already loaded AnnData object"
        )

    X, local_means, local_vars, batch_indices_, labels = (
        GeneExpressionDataset.get_attributes_from_matrix(data, labels=labels)
    )
    batch_indices = batch_indices if batch_indices is not None else batch_indices_

    super().__init__(
        X,
        local_means,
        local_vars,
        batch_indices,
        labels,
        gene_names=gene_names,
        cell_types=cell_types,
    )

    self.subsample_genes(new_n_genes=new_n_genes, subset_genes=subset_genes)


# Traceback (MemoryError)
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-3-138de299bfb8> in <module>
1 # load
2 dataset = AnnDataset('data/dataset.h5ad',
----> 3                      save_path=path)

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scvi/dataset/anndata.py in __init__(self, filename, save_path, url, new_n_genes, subset_genes)
36         self.url = url
37
---> 38         data, gene_names = self.download_and_preprocess()
39
40         super().__init__(*GeneExpressionDataset.get_attributes_from_matrix(data),

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scvi/dataset/dataset.py in download_and_preprocess(self)
61     def download_and_preprocess(self):
62         self.download()
---> 63         return self.preprocess()
64
65     def collate_fn(self, batch):

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scvi/dataset/anndata.py in preprocess(self)
52             data = ad.X.copy()  # Dense
53         else:
---> 54             data = ad.X.toarray()  # Sparse
55         select = data.sum(axis=1) > 0  # Take out cells that doesn't express any gene
56         data = data[select, :]

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scipy/sparse/compressed.py in toarray(self, order, out)
960         if out is None and order is None:
961             order = self._swap('cf')[0]
--> 962         out = self._process_toarray_args(order, out)
963         if not (out.flags.c_contiguous or out.flags.f_contiguous):
964             raise ValueError('Output array must be C or F contiguous')

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scipy/sparse/base.py in _process_toarray_args(self, order, out)
1185             return out
1186         else:
-> 1187             return np.zeros(self.shape, dtype=self.dtype, order=order)
1188
1189

MemoryError:

# Fixed code
def __init__(
    self,
    filename_or_anndata: Union[str, anndata.AnnData],
    save_path: str = "data/",
    url: str = None,
    new_n_genes: bool = False,
    subset_genes: List[int] = None,
):
    if type(filename_or_anndata) == str:
        self.download_name = filename_or_anndata
        self.save_path = save_path
        self.url = url

        data, gene_names, batch_indices, cell_types, labels = (
            self.download_and_preprocess()
        )

    elif isinstance(filename_or_anndata, anndata.AnnData):
        ad = filename_or_anndata
        data, gene_names, batch_indices, cell_types, labels = (
            self.extract_data_from_anndata(ad)
        )

    else:
        raise Exception(
            "Please provide a filename of an AnnData file or an already loaded AnnData object"
        )
    X, local_means, local_vars, batch_indices_, labels = (
        GeneExpressionDataset.get_attributes_from_matrix(data, labels=labels)
    )
    batch_indices = batch_indices if batch_indices is not None else batch_indices_

    super().__init__(
        X,
        local_means,
        local_vars,
        batch_indices,
        labels,
        gene_names=gene_names,
        cell_types=cell_types,
    )

    self.subsample_genes(new_n_genes=new_n_genes, subset_genes=subset_genes)