I need to set a boolean flag in the code running on multiple GPUs. When using a single GPU, self.calculate_running is being set to False correctly after the first iteration. It’s not being set when I use more than one GPU:
class PCTL_Layer(nn.Module):
    def __init__(self, calculate_running=False):
        super(PCTL_Layer, self).__init__()
        self.register_buffer('running', torch.zeros(1))
        self.calculate_running = calculate_running

    def forward(self, input):
        if self.calculate_running:
            pctl, _ = torch.kthvalue(input.view(-1), int(input.numel() * 0.9))
            self.running = pctl
            print(' gpu {}  calculate_running: {}'.format(torch.cuda.current_device(), self.calculate_running))
            self.calculate_running = False

        return self.running


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.percentile = PCTL_Layer(calculate_running=True)

    def forward(self, x):
        x = self.percentile(x)
        return x

model = Model()
model = torch.nn.DataParallel(model).cuda()

for i, image in enumerate(train_loader):
    output = model(image)

I tried inserting torch.cuda.synchronize() everywhere (before and after), but that didn’t help. Any advice on how to make it work?
p.s. I realize that a simple solution would be to pass i to the forward method of PCTL_layer, however it’s complicated because that layer is packaged into nn.Sequential list, so I don’t know how to do that either.