scGen model is very slow

[joseph-siefert]

Thanks for the great tools! Is the scGen model in scArches using the newest scGen version? I can run scGen no problem and it runs pretty fast, but when I run the scGen model with scArches is very slow. I would like to run with scArches so I can map query data onto the reference after correcting batch effects in the reference.

#runs fine
scgen.SCGEN.setup_anndata(adata, batch_key="dataset", labels_key="cell_type")
model = scgen.SCGEN(adata)
model.train(
    max_epochs=100,
    batch_size=32,
    early_stopping=True,
    early_stopping_patience=25,
)

#super slow
epoch = 50
early_stopping_kwargs = {
    "early_stopping_metric": "val_loss", #I have also tried elbo_metric and still very slow
    "patience": 20,
    "threshold": 0,
    "reduce_lr": True,
    "lr_patience": 13,
    "lr_factor": 0.1,
}
network = sca.models.scgen(adata = source_adata , hidden_layer_sizes=[256,128])
network.train(n_epochs=epoch, early_stopping_kwargs=early_stopping_kwargs, use_gpu=True)

Am I missing something, or is the scGen model in scArches not optimal?