fix: serialize head_dtype in all ImageClassifier subclasses

keras_hub/src/models/deit/deit_image_classifier.py

@@ -153,6 +153,7 @@ def __init__(
 
         # === config ===
         self.num_classes = num_classes
+        self.head_dtype = head_dtype
         self.pooling = pooling
         self.activation = activation
         self.dropout = dropout
@@ -166,6 +167,7 @@ def get_config(self):
                 "pooling": self.pooling,
                 "activation": self.activation,
                 "dropout": self.dropout,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config

keras_hub/src/models/hgnetv2/hgnetv2_image_classifier.py

@@ -201,6 +201,7 @@ def __init__(
         self.pooling = pooling
         self.dropout = dropout
         self.num_classes = num_classes
+        self.head_dtype = head_dtype
 
     def get_config(self):
         config = Task.get_config(self)
@@ -211,6 +212,7 @@ def get_config(self):
                 "activation": self.activation,
                 "dropout": self.dropout,
                 "head_filters": self.head_filters,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config

keras_hub/src/models/image_classifier.py

@@ -151,6 +151,7 @@ def __init__(
         self.activation = activation
         self.pooling = pooling
         self.dropout = dropout
+        self.head_dtype = head_dtype
 
     def get_config(self):
         # Backbone serialized in `super`
@@ -161,6 +162,7 @@ def get_config(self):
                 "pooling": self.pooling,
                 "activation": self.activation,
                 "dropout": self.dropout,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config

keras_hub/src/models/mobilenet/mobilenet_image_classifier.py

@@ -70,6 +70,7 @@ def __init__(
 
         # === Config ===
         self.num_classes = num_classes
+        self.head_dtype = head_dtype
         self.num_features = num_features
 
     def get_config(self):
@@ -79,6 +80,7 @@ def get_config(self):
             {
                 "num_classes": self.num_classes,
                 "num_features": self.num_features,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config

keras_hub/src/models/mobilenetv5/mobilenetv5_image_classifier.py

@@ -140,6 +140,7 @@ def __init__(
 
         # === Config ===
         self.num_classes = num_classes
+        self.head_dtype = head_dtype
         self.head_hidden_size = head_hidden_size
         self.global_pool_type = global_pool
         self.drop_rate = drop_rate
@@ -152,6 +153,7 @@ def get_config(self):
                 "head_hidden_size": self.head_hidden_size,
                 "global_pool": self.global_pool_type,
                 "drop_rate": self.drop_rate,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config

keras_hub/src/models/vgg/vgg_image_classifier.py

@@ -195,6 +195,7 @@ def __init__(
 
         # === Config ===
         self.num_classes = num_classes
+        self.head_dtype = head_dtype
         self.activation = activation
         self.pooling = pooling
         self.pooling_hidden_dim = pooling_hidden_dim
@@ -211,6 +212,7 @@ def get_config(self):
                 "activation": self.activation,
                 "pooling_hidden_dim": self.pooling_hidden_dim,
                 "dropout": self.dropout,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config

keras_hub/src/models/vit/vit_image_classifier.py

@@ -167,6 +167,7 @@ def __init__(
 
         # === config ===
         self.num_classes = num_classes
+        self.head_dtype = head_dtype
         self.pooling = pooling
         self.intermediate_dim = intermediate_dim
         self.activation = activation
@@ -182,6 +183,7 @@ def get_config(self):
                 "intermediate_dim": self.intermediate_dim,
                 "activation": self.activation,
                 "dropout": self.dropout,
+                "head_dtype": keras.dtype_policies.serialize(self.head_dtype),
             }
         )
         return config