code: def gradcheck(nbounces):
      system = BouncingBallExample()
  
      variables = {
          "init_pos": system.init_pos,
          "init_vel": system.init_vel,
          "t0": system.t0,
          "gravity": system.gravity,
          "log_radius": system.log_radius,
      }
  
      event_t = system.get_collision_times(nbounces)[-1]
      event_t.backward()
  
      analytical_grads = {}
      for name, p in system.named_parameters():
          for var in variables.keys():
              if var in name:
                  analytical_grads[var] = p.grad
  
      eps = 1e-3
  
      fd_grads = {}
  
      for var, param in variables.items():
          orig = param.data
          param.data = orig - eps
          f_meps = system.get_collision_times(nbounces)[-1]
          param.data = orig + eps
          f_peps = system.get_collision_times(nbounces)[-1]
          param.data = orig
          fd = (f_peps - f_meps) / (2 * eps)
          fd_grads[var] = fd
  
      success = True
      for var in variables.keys():
          analytical = analytical_grads[var]
          fd = fd_grads[var]
          if torch.norm(analytical - fd) > 1e-4:
              success = False
              print(
                  f"Got analytical grad {analytical.item()} for {var} param but finite difference is {fd.item()}"
              )
  
      if not success:
          raise Exception("Gradient check failed.")
  
      print("Gradient check passed.")
  
file_path: torchdiffeq/examples/bouncing_ball.py
imports: import argparse
  import matplotlib.pyplot
  import torch
  import torch.nn
  from torchdiffeq import odeint, odeint_adjoint
  from torchdiffeq import odeint_event
label_desc: Use image node_id 1 for calling a global function with example usage: gradcheck(nbounces) without return types