code: def assert_grads_ok(logits, bp_iters=None):
      def fn(logits):
          d = dist.OneTwoMatching(logits, bp_iters=bp_iters)
          return d.log_partition_function
  
      torch.autograd.gradcheck(fn, logits, atol=1e-3, rtol=1e-3)
  
imports: import logging
  import math
  import pytest
  import torch
  import pyro.distributions
  from tests.common import assert_close, assert_equal, xfail_if_not_implemented
label_desc: Use image node_id 9 for calling a global function with example usage: assert_grads_ok(logits, bp_iters) and returns: d