Why is the pruning not working on Multi-GPU?

[atrah22]

Hello,

I am using 2 V100 GPUs. I trying to apply the pruning technique on pre-trained resnet50 model for imagenet dataset.

The pruning is running. However, it is using only 1 of the two GPUs. I am running the python script in the following way (I made myself the resnet50_prune.py without changing any of the pruning code).
CUDA_VISIBLE_DEVICES=0,1 python ./run/resnet-50/resnet50_prune.py