Include Head in the weights?

[deanflaviano]

I noticed the head is not included in the weights, are there any plans on adding it? Happy holidays!

if head_type == "attention-pool":
            self.head = layers.AttentionPoolingClassifier(
                embed_dim,
                out_features=embed_dim,
                num_heads=head_num_heads,
                num_queries=head_num_queries,
                use_batch_norm=False,
                qkv_bias=False,
                linear_bias=head_linear_bias,
                average_pool=head_average_pool,
            )