efficientnetv2 |TypeError: "inner_optimizer" must be an instance of `tf.keras.optimizers.Optimizer`, but got: <tensorflow.python.training.rmsprop.RMSPropOptimizer object at 0x7f95332e2250>

[exx8]

I've tried to run efficientnetv2 on my gpu. This is a sequel to #1153 . python3 main.py --mode=train --model_name=efficientnetv2-s --dataset_cfg=imagenet --model_dir=~/effi/models/efficientnetv2 --use_tpu=False I get the following error:

I0503 13:02:08.796464 140283249366848 utils.py:411] use mixed precision policy name float32
<module 'keras.api._v2.keras.mixed_precision' from '/home/eran/.local/lib/python3.8/site-packages/keras/api/_v2/keras/mixed_precision/__init__.py'>
I0503 13:02:09.335912 140283249366848 main.py:103] backbone params/flops = 21.4585 M / 1.8090 B
I0503 13:02:09.336059 140283249366848 utils.py:339] Adding scalar summary ('model/params', 21.458488)
I0503 13:02:09.338670 140283249366848 utils.py:339] Adding scalar summary ('model/flops', 1.80903990583667)
I0503 13:02:09.364594 140283249366848 utils.py:339] Adding scalar summary ('train/weight_decay', <tf.Tensor 'mul_1:0' shape=() dtype=float32>)
I0503 13:02:09.438328 140283249366848 utils.py:339] Adding scalar summary ('loss/l2reg', <tf.Tensor 'mul_2:0' shape=() dtype=float32>)
I0503 13:02:09.439843 140283249366848 utils.py:339] Adding scalar summary ('loss/xent', <tf.Tensor 'softmax_cross_entropy_loss/value:0' shape=() dtype=float32>)
I0503 13:02:09.448067 140283249366848 utils.py:339] Adding scalar summary ('train/epoch', <tf.Tensor 'truediv_1:0' shape=() dtype=float32>)
I0503 13:02:09.456434 140283249366848 utils.py:339] Adding scalar summary ('train/lr', <tf.Tensor 'cond/Identity:0' shape=() dtype=float32>)
I0503 13:02:09.457878 140283249366848 utils.py:161] Using RMSProp optimizer
INFO:tensorflow:training_loop marked as finished
I0503 13:02:09.458009 140283249366848 error_handling.py:115] training_loop marked as finished
WARNING:tensorflow:Reraising captured error
W0503 13:02:09.458083 140283249366848 error_handling.py:149] Reraising captured error
Traceback (most recent call last):
  File "main.py", line 507, in <module>
    app.run(main)
  File "/home/eran/.local/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/eran/.local/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "main.py", line 499, in main
    est.train(
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3102, in train
    rendezvous.raise_errors()
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/error_handling.py", line 150, in raise_errors
    six.reraise(typ, value, traceback)
  File "/home/eran/.local/lib/python3.8/site-packages/six.py", line 703, in reraise
    raise value
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3092, in train
    return super(TPUEstimator, self).train(
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 360, in train
    loss = self._train_model(input_fn, hooks, saving_listeners)
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1186, in _train_model
    return self._train_model_default(input_fn, hooks, saving_listeners)
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1214, in _train_model_default
    estimator_spec = self._call_model_fn(features, labels, ModeKeys.TRAIN,
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 2933, in _call_model_fn
    return super(TPUEstimator, self)._call_model_fn(features, labels, mode,
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1174, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 3191, in _model_fn
    estimator_spec = model_fn_wrapper.call_without_tpu(
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 1714, in call_without_tpu
    return self._call_model_fn(features, labels, is_export_mode=is_export_mode)
  File "/home/eran/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 2057, in _call_model_fn
    estimator_spec = self._model_fn(features=features, **kwargs)
  File "main.py", line 174, in model_fn
    optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer)
  File "/home/eran/.local/lib/python3.8/site-packages/keras/mixed_precision/loss_scale_optimizer.py", line 405, in __init__
    raise TypeError(msg)
TypeError: "inner_optimizer" must be an instance of `tf.keras.optimizers.Optimizer`, but got: <tensorflow.python.training.rmsprop.RMSPropOptimizer object at 0x7f95332e2250>.