@ptrblck  Let me ask an additional question.
How can i choose where those estimators  are stored?
I guess they are stored in gpu0 by default but even using several gpus, the gpu which stores those estimators get out of memory. I would like to use 1 gpu for computing them and the rest for dealing with the model using dataparallel.
Is that possible in pytorch?