code: def _test_clip_l2_grad_norm_(gpu):
      if gpu >= 0:
          device = torch.device("cuda:{}".format(gpu))
      else:
          device = torch.device("cpu")
      model = nn.Sequential(
          nn.Linear(2, 10),
          nn.ReLU(),
          nn.Linear(10, 3),
      ).to(device)
      x = torch.rand(7, 2).to(device)
  
      def backward():
          model.zero_grad()
          loss = model(x).mean()
          loss.backward()
  
      backward()
      raw_grads = _get_grad_vector(model)
  
      # Threshold large enough not to affect grads
      th = 10000
      backward()
      nn.utils.clip_grad_norm_(model.parameters(), th)
      clipped_grads = _get_grad_vector(model)
  
      backward()
      clip_l2_grad_norm_(model.parameters(), th)
      our_clipped_grads = _get_grad_vector(model)
  
      np.testing.assert_allclose(raw_grads, clipped_grads)
      np.testing.assert_allclose(raw_grads, our_clipped_grads)
  
      # Threshold small enough to affect grads
      th = 1e-2
      backward()
      nn.utils.clip_grad_norm_(model.parameters(), th)
      clipped_grads = _get_grad_vector(model)
  
      backward()
      clip_l2_grad_norm_(model.parameters(), th)
      our_clipped_grads = _get_grad_vector(model)
  
      with pytest.raises(AssertionError):
          np.testing.assert_allclose(
              raw_grads, clipped_grads, rtol=1e-5
          )
  
      with pytest.raises(AssertionError):
          np.testing.assert_allclose(
              raw_grads, our_clipped_grads, rtol=1e-5
          )
  
      np.testing.assert_allclose(
          clipped_grads, our_clipped_grads, rtol=1e-5
      )
  
file_path: pfrl/tests/utils_tests/test_clip_l2_grad_norm.py
imports: import timeit
  from logging import getLogger
  import numpy
  import pytest
  import torch
  from torch import nn
  from pfrl.utils import clip_l2_grad_norm_
label_desc: Use image node_id 2 for calling a global function with example usage: _test_clip_l2_grad_norm_(gpu) without return types