Multi-GPU training: how to make some code run only one of the GPUs ?

[Zhao Kai]

Hi ALL,

I'm developing my own new layer, and I have some special operations in this layer and It works well on single GPU training.

But when It works on multi-GPU, I want my special operation only be activate on one one of the GPUs. By default, every GPU will have the same actions. How can I make it possible that some special code run only on one of the GPUs ?

Thanks.