# Buggy code
def _get_batch_code_from_category(
    adata: anndata.AnnData, category: Sequence[Union[int, str]]
):
    categorical_mappings = adata.uns["_scvi"]["categorical_mappings"]
    batch_mappings = categorical_mappings["_scvi_batch"]["mapping"]
    batch_code = []
    for cat in category:
        if cat is None:
            batch_code.append(None)
        elif cat not in batch_mappings:
            raise ValueError('"{}" not a valid batch category.'.format(cat))
        else:
            batch_loc = np.where(batch_mappings == cat)[0][0]
            batch_code.append(batch_loc)
    return batch_code


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-143-b1df721788f5> in <module>
2 for tissue in np.unique(adata.obs['tissue']):
3     sub_adata = adata[adata.obs['tissue']==tissue]
----> 4     de_celltype[tissue] = vae.differential_expression(sub_adata, groupby = 'Propagated.Annotation', batch_correction=True)

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/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)
187             batch_size=batch_size,
188         )
--> 189         result = _de_core(
190             adata,
191             model_fn,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/_utils.py in _de_core(adata, model_fn, groupby, group1, group2, idx1, idx2, all_stats, all_stats_fn, col_names, mode, batchid1, batchid2, delta, batch_correction, **kwargs)
59             cell_idx2 = (adata.obs[groupby] == group2).to_numpy().ravel()
60
---> 61         all_info = dc.get_bayes_factors(
62             cell_idx1,
63             cell_idx2,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/utils/differential.py in get_bayes_factors(self, idx1, idx2, mode, batchid1, batchid2, use_observed_batches, n_samples, use_permutation, m_permutation, change_fn, m1_domain_fn, delta, cred_interval_lvls)
168         eps = 1e-8  # used for numerical stability
169         # Normalized means sampling for both populations
--> 170         scales_batches_1 = self.scale_sampler(
171             selection=idx1,
172             batchid=batchid1,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_no_grad(*args, **kwargs)
47         def decorate_no_grad(*args, **kwargs):
48             with self:
---> 49                 return func(*args, **kwargs)
50         return decorate_no_grad
51

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/utils/differential.py in scale_sampler(self, selection, n_samples, n_samples_per_cell, batchid, use_observed_batches, give_mean)
392             idx = np.random.choice(np.arange(self.adata.shape[0])[selection], n_samples)
393             px_scales.append(
--> 394                 self.model_fn(self.adata, indices=idx, transform_batch=batch_idx)
395             )
396             batch_idx = batch_idx if batch_idx is not None else np.nan

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_no_grad(*args, **kwargs)
47         def decorate_no_grad(*args, **kwargs):
48             with self:
---> 49                 return func(*args, **kwargs)
50         return decorate_no_grad
51

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/rnamixin.py in get_normalized_expression(self, adata, indices, transform_batch, gene_list, library_size, n_samples, batch_size, return_mean, return_numpy)
82         scdl = self._make_scvi_dl(adata=adata, indices=indices, batch_size=batch_size)
83         if transform_batch is not None:
---> 84             transform_batch = _get_batch_code_from_category(adata, transform_batch)
85
86         if gene_list is None:

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/model/_utils.py in _get_batch_code_from_category(adata, category)
131     batch_mappings = categorical_mappings["_scvi_batch"]["mapping"]
132     if category not in batch_mappings:
--> 133         raise ValueError('"{}" not a valid batch category.'.format(category))
134     return np.where(batch_mappings == category)[0][0]

ValueError: "0" not a valid batch category.```

#### Versions:
<!-- Output of scvi.__version__ -->
VERSION

<!-- Relevant screenshots -->

# Fixed code
def _get_batch_code_from_category(
    adata: anndata.AnnData, category: Sequence[Union[Number, str]]
):
    if not isinstance(category, IterableClass) or isinstance(category, str):
        category = [category]

    categorical_mappings = adata.uns["_scvi"]["categorical_mappings"]
    batch_mappings = categorical_mappings["_scvi_batch"]["mapping"]
    batch_code = []
    for cat in category:
        if cat is None:
            batch_code.append(None)
        elif cat not in batch_mappings:
            raise ValueError('"{}" not a valid batch category.'.format(cat))
        else:
            batch_loc = np.where(batch_mappings == cat)[0][0]
            batch_code.append(batch_loc)
    return batch_code