Hi, all
When training scVI models, I wanted to change the device used. So, I set vae.to_device(‘cuda:1’) to use another gpu. However, it didn’t work as the warning info showed LOCALRNAK:0. How could I use other divices for training?
Hi,
In the train function, you can specify the device id you wish to train the model on. For example:
model.train(
max_epochs=100,
train_size=0.5,
check_val_every_n_epoch=1,
accelerator="gpu",
devices=[1],
#strategy="ddp_find_unused_parameters_true", #for multiGPU
)
Will run it on gpu #1.
Regardless, in scvi-tools v1.3.0 we now have the option to train models in multiGPU (just use the strategy parameter like in the comment above, but then remove the devices), if that is relevant for you ,give it a try!
1 Like
Thanks, it does work!