code: def test_grad_norm_with_double_dtype():
      class Model(nn.Module):
          def __init__(self):
              super().__init__()
              dtype = torch.double
              self.param = nn.Parameter(torch.tensor(1.0, dtype=dtype))
              # grad norm of this would become infinite
              self.param.grad = torch.tensor(1e23, dtype=dtype)
  
      model = Model()
      norms = grad_norm(model, 2)
      assert all(
          torch.isfinite(torch.tensor(v)) for v in norms.values()
      ), norms
  
file_path: pytorch-lightning/tests/tests_pytorch/utilities/test_grads.py
imports: from unittest.mock import Mock
  import pytest
  import torch
  import torch.nn
  from lightning.pytorch.utilities import grad_norm
label_desc: Use image node_id 3 for calling a global function with example usage: test_grad_norm_with_double_dtype() without return types