code: def create_default_gradient_fn(forward_fn, loss_outputs, loss_fn):
      """
      This function calls the gradient function, to implement the backpropogation
      """
  
      @jax.jit
      def model_loss(params, batch, target, weights, rng):
          predict = forward_fn(params, rng, batch)
          if loss_outputs is not None:
              predict = [predict[i] for i in loss_outputs]
          return loss_fn(predict, target, weights)
  
      return model_loss
  
imports: import numpy
  import time
  import logging
  from collections.abc import Sequence
  from deepchem.data import Dataset, NumpyDataset
  from deepchem.metrics import Metric
  from deepchem.models.models import Model
  from deepchem.models.losses import Loss
  from deepchem.models.optimizers import Optimizer, Adam
  from deepchem.utils.evaluate import GeneratorEvaluator
  from deepchem.trans.transformers import Transformer, undo_transforms
  from typing import Any, Callable, Iterable, List, Optional, Tuple, Union, Sequence
  from deepchem.utils.typing import LossFn, OneOrMany
  import jax.numpy
  import jax
  import haiku
  import optax
  import warnings
label_desc: Use image node_id 3 for calling a global function with example usage: create_default_gradient_fn(forward_fn, loss_outputs, loss_fn) and returns: model_loss, loss_fn