Diff Coverage

            local_p = torch.linalg.vector_norm(
                torch.stack(norms).to(torch.float32), norm_type,
            ) ** norm_type
        else:
            local_p = torch.tensor(0.0, device=device, dtype=torch.float32)

        for group in sig_groups[sig]:
            dist.all_reduce(local_p, op=dist.ReduceOp.SUM, group=group)