def METHOD_NAME(self):
    shared_device = self.param_groups[0]["params"][
        0
    ].device  # put everything on the same device, in case of model parallelism
    norm = torch.norm(
        torch.stack(
            [
                ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for group in self.param_groups
                for p in group["params"]
                if p.grad is not None
            ]
        ),
        p=2,
    )
    return norm