code: def test_multi_dim_weighted_train(self):
      with tf.Graph().as_default():
          weights_feature_name = self._default_weights_feature_name
  
          def _train_op_fn(loss):
              return loss
  
          head = ranking_head.create_ranking_head(
              loss_fn=_make_loss_fn(weights_feature_name),
              train_op_fn=_train_op_fn,
          )
          # Create estimator spec.
          spec = head.create_estimator_spec(
              features={weights_feature_name: self._default_weights},
              mode=tf_estimator.ModeKeys.TRAIN,
              logits=self._default_logits,
              labels=self._default_labels,
          )
  
          # Assert predictions, loss, and train_op.
          with self.cached_session() as sess:
              _initialize_variables(self, spec.scaffold)
              loss, train_result = sess.run((spec.loss, spec.train_op))
              self.assertAllClose(self._default_weighted_loss, loss)
              self.assertAllClose(
                  self._default_weighted_loss, train_result
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tensorflow import estimator
  from tensorflow_ranking.python import head
  from tensorflow_ranking.python import metrics
label_desc: Use image node_id 9 for calling the RankingHeadTest obj's underlying member method code with example usage: obj.test_multi_dim_weighted_train() and returns: loss