I don’t think that “measure_module_sparsity()” function will get called on a module such as “ReLU” because within “measure_global_sparsity()” function, the for loop:
for module_name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        module_num_zeros, module_num_elements, _ = measure_module_sparsity(
                module, weight=weight, bias=bias, use_mask=conv2d_use_mask)
        num_zeros += module_num_zeros
        num_elements += module_num_elements

will make sure that “measure_module_sparsity()” is only called for Conv2d and/or Linear layers.
At least, that’s what my understanding of the code is.