Can not squeeze dim[1], expected a dimension of 1, got 10 (Masking Layer)

Question

Given the MWE below, Python=3.8.15, Tensorflow=2.11, I get the error as shown. Note that this does not occur when layers.Masking(mask_value=0.0, name='mask1') is commented out.

import tensorflow as tf
import numpy as np
from tensorflow.keras import layers

samples, timesteps, features = 32, 10, 8
inputs = np.random.random([samples, timesteps, features]).astype(np.float32)
inputs[:, 3, :] = 0.
inputs[:, 5, :] = 0.
print(inputs.shape)

labels = np.mean(inputs, axis=2)
print(labels.shape)

model = tf.keras.models.Sequential([
    layers.Input(shape=(timesteps, features), name='input1'),
    layers.Masking(mask_value=0.0, name='mask1'),
    layers.Bidirectional(layers.LSTM(32, return_sequences=True, name='lstm1'), name='bilstm1'),
    layers.Dropout(0.2, name='dropout1'),
    layers.Dense(1, activation='linear', name='dense2')
])

output = model(inputs)
print(output.shape)

model.compile(
    loss        = tf.keras.losses.mean_squared_error,
    optimizer   = tf.keras.optimizers.Adam(),
    metrics     = [
        tf.keras.metrics.mean_absolute_error,
    ],
    run_eagerly = False,
)

print(model.summary())

history = model.fit(
    x      = inputs,
    y      = labels,
    epochs = 30
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], line 36
     25 model.compile(
     26     loss        = tf.keras.losses.mean_squared_error,
     27     optimizer   = tf.keras.optimizers.Adam(),
   (...)
     31     run_eagerly = False,
     32 )
     34 print(model.summary())
---> 36 history = model.fit(
     37     x      = inputs,
     38     y      = labels,
     39     epochs = 30
     40 )
...
    File "/home/dve/anaconda3/envs/trt/lib/python3.8/site-packages/keras/utils/losses_utils.py", line 224, in squeeze_or_expand_dimensions
        sample_weight = tf.squeeze(sample_weight, [-1])

    ValueError: Can not squeeze dim[1], expected a dimension of 1, got 10 for '{{node mean_squared_error/weighted_loss/Squeeze}} = Squeeze[T=DT_FLOAT, squeeze_dims=[-1]](mean_squared_error/mul_1)' with input shapes: [32,10].

Answer 1

Your labels have the wrong shape of (32,10) but TF requires (32, 10, 1). It tries to squeeze the last dimension (which it expects to be 1) but gets 10.

You can fix it by keeping the dimension when calculating the labels as follows:

labels = np.mean(inputs, axis=2, keepdims=True)