# Generate predictions
y_pred_prob = model.predict(normalized_test_dataset)  # Predicted probabilities
y_pred = np.argmax(y_pred_prob, axis=1)  # Convert probabilities to class indices

# Extract true labels from dataset and convert to class indices
y_true = np.concatenate([np.argmax(label.numpy(), axis=1) for _, label in normalized_test_dataset])
print(f"y_true shape: {y_true.shape}, unique values: {np.unique(y_true)}")
print(f"y_pred shape: {y_pred.shape}, unique values: {np.unique(y_pred)}")

num_uniq,count = np.unique(y_pred,return_counts=True)
print(num_uniq,count)

num_uniq,count = np.unique(y_true,return_counts=True)
print(num_uniq,count)

# Validate shapes
print(f"y_pred shape: {y_pred.shape}, y_true shape: {y_true.shape}")

        __
        __