Hi, I try to register hook and run it on multiple GPUs. However, it only return the result on GPU 0.
The data has been successfully split on multiple GPUs.
Anyone know why is it?
The code:
def forward(self, x):
        self.activations = []
        self.gradients = []
        self.grad_index = 0
        self.activation_to_layer = {}

        activation_index = 0

        for layer, module in self.model.named_modules():
            if ('conv' in layer) or ('pool' in layer) or ('fc' in layer):
                if 'fc6' in layer:
                    if isinstance(self.model, nn.DataParallel):
                        x = x.view(-1, self.model.module.fc6.in_features)
                    else:
                        x = x.view(-1, self.model.fc6.in_features)
                x = module(x)
            if isinstance(module, torch.nn.modules.conv.Conv3d):
                # hook will registered on the output of the layer
                x.register_hook(self.compute_rank)
                self.activations.append(x)
                self.activation_to_layer[activation_index] = layer
                activation_index += 1
                x = model.relu(x)
            elif isinstance(module, torch.nn.modules.Linear) and layer != 'fc8':
                x = model.dropout(model.relu(x))

        return x

def compute_rank(self, grad):
        """
        Compute the Taylor expansion without abs of each channel
        return:
            self.activations: feature map before relu in each layer
            self.filter_ranks: Taylor value without abs over spatial and batch
        """
        activation_index = len(self.activations) - self.grad_index - 1
        activation = self.activations[activation_index]
        if self.pruning_level == 'channel':
            print(self.model)
            print('device', grad.device)
            values = torch.sum((activation * grad), dim = 4).\
                        sum(dim=3).sum(dim=2).sum(dim=0)

            # Normalize the rank by the filter dimensions
            values = values / (activation.size(0) * activation.size(2) \
                                * activation.size(3) * activation.size(4))

            if activation_index not in self.filter_ranks:
                self.filter_ranks[activation_index] = \
                    torch.FloatTensor(activation.size(1)).zero_().cuda()

and the result:
device cuda:0