Check Net Device

[ali-akhavan89]

Hi all. I haven't gone fully into the libraries, but I've noticed that in factory.py under posterior_nn there's a hard-coded 'cpu' for checking the net device:
check_net_device(embedding_net, "cpu", embedding_net_warn_msg)
Does this mean that when my desired device is CUDA, I should also change this? I'm asking this because I've been getting this message lately in one of our code bases:
UserWarning: The passed embedding net will be moved to cpu for constructing the net building function. check_net_device(embedding_net, "cpu", embedding_net_warn_msg),

thank you,
Ali

[janfb]

Hi Ali, this happens during composition of the user-specified density estimator ("mdn", "maf" etc), the z-scoring transforms, and the potential user-specified embedding net. The composition of the these nets always happens on the cpu, therefore the hardcoded "cpu".

Thus, if you are passing an embedding net to posterior_nn then this net should be on the CPU, even if you later want to train on GPU.

This is how it could look like:

net_builder = posterior_nn(
    model="maf", 
    z_score_x="structured", 
    embedding_net=your_embedding_on_cpu
)

trainer = NPE(prior=prior, density_estimator=net_builder, device="cuda")

passing device="cuda" (or your local device name), is all you need to train on GPU.

Does this answer your question?

[ali-akhavan89]

Jan, thank you so much for your answer and such a clear instruction.
Best,
Ali