tf_mtl_single_task_head.py

def single_head_cnn_model_fn(features, labels, mode):
  
  # Extract the features
  dense = extract_features(features)
  
  # Predictions
  predictions = tf.layers.dense(inputs=dense, units=2)

  # Optimizer
  optimizer = tf.train.AdamOptimizer()
  
  # Define the head
  regression_head = tf.contrib.estimator.regression_head(label_dimension=2)
  return regression_head.create_estimator_spec(features, mode, predictions, labels[:, 2:8:5], optimizer)