Sample working code which computes intermediate layer output and cosine similarity loss

import tensorflow as tf
(x_train, y_train), (x_test, y_test)=tf.keras.datasets.mnist.load_data()
model = tf.keras.Sequential([
tf.keras.layers.Flatten(input_shape=[28,28,1]),
tf.keras.layers.Dense(64,activation=“relu”),
tf.keras.layers.Dense(10, activation=“softmax”)
])
intermediate_layer = tf.keras.Model(inputs=model.inputs,
outputs=model.outputs + [model.layers[1].output])
cosine_loss = tf.keras.losses.CosineSimilarity()
optimezer = tf.optimizers.Adam()
with tf.GradientTape() as tape:
y_pred, intermediate_output= intermediate_layer(x_train)
print(intermediate_output)
loss = cosine_loss(tf.one_hot(y_train, 10), y_pred)
gradients = tape.gradient(loss, intermediate_layer.variables)

Output

tf.Tensor(
[[1.2817659e+02 0.0000000e+00 2.1529781e+01 … 1.4793573e+02
1.7081753e+01 0.0000000e+00]
[3.1562555e+02 0.0000000e+00 8.2275620e+01 … 0.0000000e+00
0.0000000e+00 0.0000000e+00]
[1.5004303e+01 5.3053143e+01 0.0000000e+00 … 0.0000000e+00
0.0000000e+00 9.4606926e+01]
…
[0.0000000e+00 0.0000000e+00 1.2034170e+02 … 1.3093190e+02
0.0000000e+00 1.5922928e-01]
[2.4972404e+02 0.0000000e+00 2.2767758e+02 … 0.0000000e+00
0.0000000e+00 2.0227030e+01]
[9.6311852e+01 0.0000000e+00 0.0000000e+00 … 5.8486767e+00
0.0000000e+00 9.8323425e+01]], shape=(60000, 64), dtype=float32)