batch_removal AssertionError

Hi, I have a similar issue to #30 with batch_removal. The training of the network works, but running batch_removal throws an error:
adata = sc.read(path+"Aggr.adata.h5ad", cache=True)
network = scgen.VAEArithKeras(x_dimension= adata.shape[1], model_path="./models/batch")
network.train(train_data=adata, n_epochs=100)
corrected_adata =  scgen.batch_removal(network=network, adata=adata, batch_key="batch", cell_label_key="cell_type")

 ---------------------------------------------------------------------------
 AssertionError                            Traceback (most recent call last)
<ipython-input-20-83ab6f5c386a> in <module>
      1 t = time.process_time()
----> 2 corrected_adata =  scgen.batch_removal(network=network, adata=adata, batch_key="batch", cell_label_key="cell_type")
      3 elapsed_time = time.process_time() - t

~/anaconda3/lib/python3.7/site-packages/scgen/models/util.py in batch_removal(network, adata, batch_key, cell_label_key)
    287             temp_cell[batch_ind[study]].X = batch_list[study].X
    288         shared_ct.append(temp_cell)
--> 289     all_shared_ann = anndata.AnnData.concatenate(*shared_ct, batch_key="concat_batch", index_unique=None)
    290     if "concat_batch" in all_shared_ann.obs.columns:
    291         del all_shared_ann.obs["concat_batch"]

~/anaconda3/lib/python3.7/site-packages/anndata/_core/anndata.py in concatenate(self, join, batch_key, batch_categories, uns_merge, index_unique, fill_value, *adatas)
   1703             uns_merge=uns_merge,
   1704             fill_value=fill_value,
-> 1705             index_unique=index_unique,
   1706         )
   1707 

~/anaconda3/lib/python3.7/site-packages/anndata/_core/merge.py in concat(adatas, join, batch_key, batch_categories, uns_merge, index_unique, fill_value)
    470         # Current behaviour is mostly for backwards compat. It's like make_names_unique, but
    471         # unfortunately the behaviour is different.
--> 472         partial(merge_outer, batch_keys=batch_categories, merge=merge_same),
    473     )
    474 

~/anaconda3/lib/python3.7/site-packages/anndata/_core/merge.py in merge_dataframes(dfs, new_index, merge_strategy)
    403 
    404 def merge_dataframes(dfs, new_index, merge_strategy=merge_unique):
--> 405     dfs = [df.reindex(index=new_index) for df in dfs]
    406     # New dataframe with all shared data
    407     new_df = pd.DataFrame(merge_strategy(dfs), index=new_index)

~/anaconda3/lib/python3.7/site-packages/anndata/_core/merge.py in <listcomp>(.0)
    403 
    404 def merge_dataframes(dfs, new_index, merge_strategy=merge_unique):
--> 405     dfs = [df.reindex(index=new_index) for df in dfs]
    406     # New dataframe with all shared data
    407     new_df = pd.DataFrame(merge_strategy(dfs), index=new_index)

~/anaconda3/lib/python3.7/site-packages/pandas/util/_decorators.py in wrapper(*args, **kwargs)
    225         @wraps(func)
    226         def wrapper(*args, **kwargs) -> Callable[..., Any]:
--> 227             return func(*args, **kwargs)
    228 
    229         kind = inspect.Parameter.POSITIONAL_OR_KEYWORD

~/anaconda3/lib/python3.7/site-packages/pandas/core/frame.py in reindex(self, *args, **kwargs)
   3854         kwargs.pop("axis", None)
   3855         kwargs.pop("labels", None)
-> 3856         return self._ensure_type(super().reindex(**kwargs))
   3857 
   3858     def drop(

~/anaconda3/lib/python3.7/site-packages/pandas/core/base.py in _ensure_type(self, obj)
     91         Used by type checkers.
     92         """
---> 93         assert isinstance(obj, type(self)), type(obj)
     94         return obj
     95 

AssertionError: <class 'pandas.core.frame.DataFrame'>
I will be very grateful for some help!
anndata==0.7.3 scanpy==1.5.2.dev7+ge33a2f33 scgen===1.1.5.dev2-3004bc0 scipy==1.4.1 scikit-learn==0.22.1 pandas==1.0.1
theislab / scgen

batch_removal AssertionError #32
