AttributeError: 'NoneType' object has no attribute 'replace' in model.save()

[davidwincent]

Models with the Time2Vec layer can not be saved using model.save() method

Example:

model = time2vec_lstm(SEQ_LEN, 16)

model.compile(
    loss=tf.keras.losses.mean_squared_error,
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3)
)

# (Re)fit
history = model.fit(
    x=[tr_X1.reshape((-1, SEQ_LEN, 1)), tr_X2.reshape((-1, SEQ_LEN, 1))],
    y=tr_y,
    batch_size=64, epochs=1, verbose=2,
    callbacks=[ModelCallback()]
)

# Save model
model.save('Time2Vec/models/time2vec_lstm')

Gives the following error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_10923/1350189499.py in <module>
----> 1 model.save('Time2Vec/models/time2vec_lstm')

~/anaconda3/lib/python3.9/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     65     except Exception as e:  # pylint: disable=broad-except
     66       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67       raise e.with_traceback(filtered_tb) from None
     68     finally:
     69       del filtered_tb

~/anaconda3/lib/python3.9/site-packages/tensorflow/python/training/tracking/graph_view.py in _escape_local_name(name)
     60   # edges traversed to reach the variable, so we escape forward slashes in
     61   # names.
---> 62   return (name.replace(_ESCAPE_CHAR, _ESCAPE_CHAR + _ESCAPE_CHAR)
     63           .replace(r"/", _ESCAPE_CHAR + "S"))
     64 

AttributeError: 'NoneType' object has no attribute 'replace'