load_model doesn’t work when using RSquare from tensorflow_addons as metric

I have a model that uses R2 as a metric. Since AFAIK there isn’t one natively implemented in TF, I use the one from the tensorflow-addons package. However, when I try to load this model after saving, it fails with the error:

type of argument “y_shape” must be a tuple; got list instead

Here is a minimal working example that produces this error:

from tensorflow.keras.models import load_model, Sequential from tensorflow.keras.layers import Dense, Input import tensorflow as tf import tensorflow_addons as tfa model = Sequential() model.add(Input(5)) model.add(Dense(5)) model.add(Dense(5)) model.compile(metrics = [tfa.metrics.RSquare(y_shape=(5,))]) model.save('test_model.h5') model = load_model('test_model.h5') 

RSquare works fine during training but I need to be able to load the model later (and load models I have already saved). I have tried using the custom_objects argument to load_model but this makes no difference. Any suggestions?

Thanks in advance!