Update pages/15_TransferLearning_HF.py

pages/15_TransferLearning_HF.py

@@ -3,7 +3,6 @@ import tensorflow as tf
 from transformers import ViTFeatureExtractor, TFAutoModelForImageClassification
 import tensorflow_datasets as tfds
 import matplotlib.pyplot as plt
-import numpy as np
 
 # Load the dataset
 dataset_name = "cats_vs_dogs"
@@ -35,8 +34,9 @@ base_model.trainable = False
 
 # Add custom layers on top
 inputs = tf.keras.Input(shape=(224, 224, 3))
-x = feature_extractor(inputs)
-x = tf.keras.layers.Flatten()(base_model(x)[0])
+features = feature_extractor(inputs, return_tensors="tf")["pixel_values"]
+base_output = base_model(features)[0]  # Extract base model output
+x = tf.keras.layers.Flatten()(base_output)
 x = tf.keras.layers.Dense(256, activation='relu')(x)
 x = tf.keras.layers.Dropout(0.5)(x)
 outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)