Add the output projects layer in the checkpoint items of MultiClsHead.

PiperOrigin-RevId: 361916017

official/nlp/modeling/layers/cls_head.py

@@ -158,5 +158,6 @@ class MultiClsHeads(tf.keras.layers.Layer):
 
   @property
   def checkpoint_items(self):
-    # TODO(hongkuny): add output projects to the checkpoint items.
-    return {self.dense.name: self.dense}
+    items = {self.dense.name: self.dense}
+    items.update({v.name: v for v in self.out_projs})
+    return items

official/nlp/modeling/layers/cls_head_test.py

@@ -48,7 +48,7 @@ class MultiClsHeadsTest(tf.test.TestCase):
     self.assertAllClose(outputs["foo"], [[0., 0.], [0., 0.]])
     self.assertAllClose(outputs["bar"], [[0., 0., 0.], [0., 0., 0.]])
     self.assertSameElements(test_layer.checkpoint_items.keys(),
-                            ["pooler_dense"])
+                            ["pooler_dense", "foo", "bar"])
 
   def test_layer_serialization(self):
     cls_list = [("foo", 2), ("bar", 3)]