Unable to load saved model with TPU using mixed precision bfloat16
Hello, @rcalonso!
Could you try to use the tf.keras.experimental.load_from_saved_model function with the policy="infer_float32_vars" argument. This will load the model with the appropriate mixed precision policy for inference. Also could you try with latest TF v2.15 or nightly and let us know?
Thank you! Hi @sushreebarsa, thanks for the response.

The function `tf.keras.experimental.load_from_saved_model` is not available in TF 2.14.1 and 2.15. According to https://github.com/tensorflow/tensorflow/blob/v2.15.0/tensorflow/python/keras/saving/saved_model_experimental.py#L407-L409 the function `tf.keras.experimental.load_from_saved_model` is deprecated.

Also, using the policy `infer_float32_vars` is not possible because is deprecated.
```bash
ValueError: Policies ending in '_float32_vars' have been removed from TensorFlow. Please use the 'mixed_float16' or 'mixed_bfloat16' policy instead. Got policy name: 'infer_float32_vars'
```

I tried the code I provided in the description with TF 2.15 in cloud TPU and the same error was raised. @rcalonso Thanks for replying here.
It seems that there is incompatibility between the saved model format and the current execution environment. This usually happens when a model trained with mixed precision and TPUs is attempted to be loaded on a CPU or without setting the appropriate mixed-precision settings. Install or update the tensorflow-model-optimization library if needed

Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. @sushreebarsa I'm saving the model in `.keras` format because when saving with the saved model format (`tf`) the custom train step was not saved. I created an issue for this https://github.com/tensorflow/tensorflow/issues/62450

What would be the appropriate mixed-precision settings? I've tried loading the model with float32 policy but the same error was raised. The tensorflow-model-optimization version used is the latest available. HI @rcalonso ,

Apologies for the delay. I have modified the code to Keras3 compatible and its executing fine.Please refer the attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/a331152fd4c11a6681a62b9d6702c397/62564_tpu_keras3.ipynb). Could you please check and confirm ?

Thanks! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62564">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62564">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.14.1

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I'm trying to train a custom model with TPU using mixed precision bfloat16. I'm able to train the model and save checkpoints in `.keras` format during training. However, when trying to load a saved model with `tf.keras.models.load_model` I'm getting an error which I'm not able to figure out.  

The custom model I'm trying to train uses Efficientnet as backbone and the error is raised when trying to load the mean and variance from the normalization layer of the Efficientnet model. I've reproduced the issue using only the Efficientnet model.

Using full precision, the error doesn't happen. Also, when training with mixed precision float16 using a single gpu the error doesn't happen either.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

tf.keras.mixed_precision.set_global_policy("mixed_bfloat16")

ds = (
    tf.data.Dataset.from_tensor_slices(
        (
            tf.ones([1, 224, 224, 3], dtype=tf.bfloat16),
            tf.ones([1, 1000], dtype=tf.float32),
        )
    )
    .repeat()
    .batch(batch_size=1)
)


resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu="local",
)
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
print(f"All TPU devices: {tf.config.list_logical_devices('TPU')}")

with strategy.scope():
    model = tf.keras.applications.EfficientNetB0()
    model.compile(loss="binary_crossentropy")

model.fit(
    ds,
    steps_per_epoch=1,
    epochs=1,
    callbacks=[
        tf.keras.callbacks.ModelCheckpoint(
            filepath="epoch_{epoch:03d}.keras",
            monitor="loss",
            verbose=1,
            save_weights_only=False,
            save_freq="epoch",
            options=None,
        )
    ]
)

print("Loading model")
with strategy.scope():
    loaded_model = tf.keras.models.load_model("epoch_001.keras")

print("Continue training...")
model.fit(
    ds,
    steps_per_epoch=1,
    epochs=1,
    callbacks=[
        tf.keras.callbacks.ModelCheckpoint(
            filepath="epoch_{epoch:03d}.keras",
            monitor="loss",
            verbose=1,
            save_weights_only=False,
            save_freq="epoch",
            options=None,
        )
    ]
)
```


### Relevant log output

```shell
D1204 08:52:31.331000367   42266 lb_policy_registry.cc:47]             registering LB policy factory for "cds_experimental"
D1204 08:52:31.331008502   42266 lb_policy_registry.cc:47]             registering LB policy factory for "xds_cluster_resolver_experimental"
D1204 08:52:31.331016540   42266 lb_policy_registry.cc:47]             registering LB policy factory for "xds_override_host_experimental"
D1204 08:52:31.331024721   42266 lb_policy_registry.cc:47]             registering LB policy factory for "xds_wrr_locality_experimental"
D1204 08:52:31.331033024   42266 lb_policy_registry.cc:47]             registering LB policy factory for "ring_hash_experimental"
D1204 08:52:31.331041050   42266 certificate_provider_registry.cc:33]  registering certificate provider factory for "file_watcher"
I1204 08:52:31.331790894   42266 ev_epoll1_linux.cc:360]               grpc epoll fd: 12
I1204 08:52:31.355651140   42266 socket_utils_common_posix.cc:366]     TCP_USER_TIMEOUT is available. TCP_USER_TIMEOUT will be used thereafter
2023-12-04 08:52:31.444721: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-12-04 08:52:37.879040: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x55c5b9887040 initialized for platform TPU (this does not guarantee that XLA will be used). Devices:
2023-12-04 08:52:37.879099: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): TPU, 2a886c8
2023-12-04 08:52:37.879114: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (1): TPU, 2a886c8
2023-12-04 08:52:37.879127: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (2): TPU, 2a886c8
2023-12-04 08:52:37.879140: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (3): TPU, 2a886c8
2023-12-04 08:52:37.879153: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (4): TPU, 2a886c8
2023-12-04 08:52:37.879166: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (5): TPU, 2a886c8
2023-12-04 08:52:37.879179: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (6): TPU, 2a886c8
2023-12-04 08:52:37.879192: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (7): TPU, 2a886c8
2023-12-04 08:52:37.879339: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879424: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879493: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879562: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879636: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879830: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879898: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.879957: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880013: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880085: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880270: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880349: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880407: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880468: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880533: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880711: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880778: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880845: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880917: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.880996: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881181: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881243: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881307: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881374: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881447: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881635: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881702: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881758: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881822: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.881887: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882085: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882163: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882223: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882287: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882347: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882567: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882641: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882716: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882800: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:37.882864: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:42.574430: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-12-04 08:52:42.615217: I ./tensorflow/compiler/jit/device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2023-12-04 08:52:42.641545: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:52:42.940869: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.898837: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.899128: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.899419: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.899711: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.899955: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.900451: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:00.900836: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:43.886267: I tensorflow/core/tpu/graph_rewrite/encapsulate_tpu_computations_pass.cc:241] Subgraph fingerprint:16085197774278075878
2023-12-04 08:53:45.781615: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:961] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.
2023-12-04 08:53:48.820860: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:53:48.822228: I tensorflow/core/tpu/kernels/tpu_compilation_cache_interface.cc:441] TPU host compilation cache miss: cache_key(590d039eb1480a5c:0:0), session_name()
2023-12-04 08:54:54.362634: I tensorflow/core/tpu/kernels/tpu_compile_op_common.cc:226] Compilation of 590d039eb1480a5c:0:0 with session name  took 1m5.54031078s and succeeded
2023-12-04 08:54:54.567413: I tensorflow/core/tpu/kernels/tpu_compilation_cache_interface.cc:475] TPU host compilation cache: compilation complete for cache_key(590d039eb1480a5c:0:0), session_name(), subgraph_key(std::string(property.function_name) = "cluster_train_function_16085197774278075878", property.function_library_fingerprint = 6042109918039831297, property.mlir_module_fingerprint = 0, property.num_replicas = 8, topology.chip_bounds().x = 2, topology.chip_bounds().y = 2, topology.chip_bounds().z = 1, topology.wrap().x = false, topology.wrap().y = false, topology.wrap().z = false, std::string(property.shapes_prefix) = "2,224,224,3,;2,1000,;", property.guaranteed_constants_size = 0, embedding_partitions_fingerprint = "1688352644216761960")
2023-12-04 08:54:54.567501: I tensorflow/core/tpu/kernels/tpu_compilation_cache_interface.cc:541] After adding entry for key 590d039eb1480a5c:0:0 with session_name  cache is 1 entries (112491064 bytes),  marked for eviction 0 entries (0 bytes).
2023-12-04 08:54:54.580690: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.580856: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.580995: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.581069: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.581148: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.581236: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.581315: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-04 08:54:54.581378: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
All TPU devices: [LogicalDevice(name='/device:TPU:0', device_type='TPU'), LogicalDevice(name='/device:TPU:1', device_type='TPU'), LogicalDevice(name='/device:TPU:2', device_type='TPU'), LogicalDevice(name='/device:TPU:3', device_type='TPU'), LogicalDevice(name='/device:TPU:4', device_type='TPU'), LogicalDevice(name='/device:TPU:5', device_type='TPU'), LogicalDevice(name='/device:TPU:6', device_type='TPU'), LogicalDevice(name='/device:TPU:7', device_type='TPU')]
1/1 [==============================] - ETA: 0s - loss: 0.7109
Epoch 1: saving model to epoch_001.keras
1/1 [==============================] - 101s 101s/step - loss: 0.7109
Loading model
Traceback (most recent call last):
  File "/home/ricardo_alonso/tpu_mixed_precision.py", line 48, in <module>
    loaded_model = tf.keras.models.load_model("epoch_001.keras")
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/saving/saving_api.py", line 254, in load_model
    return saving_lib.load_model(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/saving/saving_lib.py", line 281, in load_model
    raise e
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/saving/saving_lib.py", line 269, in load_model
    _load_state(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/saving/saving_lib.py", line 466, in _load_state
    _load_container_state(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/saving/saving_lib.py", line 534, in _load_container_state
    _load_state(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/saving/saving_lib.py", line 435, in _load_state
    trackable.load_own_variables(weights_store.get(inner_path))
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/layers/preprocessing/normalization.py", line 392, in load_own_variables
    super().load_own_variables(store)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/engine/base_layer.py", line 3539, in load_own_variables
    v.assign(store[f"{i}"])
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/mixed_precision/autocast_variable.py", line 297, in assign
    return self._apply_assign_update(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/keras/src/mixed_precision/autocast_variable.py", line 269, in _apply_assign_update
    assign_op = update_fn(value, use_locking, name, False)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/tpu_values.py", line 167, in assign
    return self._policy.assign(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/tpu_values.py", line 488, in assign
    return assign(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/tpu_values.py", line 427, in assign
    return var._update(  # pylint: disable=protected-access
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/values.py", line 1024, in _update
    return self._update_cross_replica(update_fn, value, **kwargs)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/values.py", line 973, in _update_cross_replica
    return self.distribute_strategy.extended.update(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/distribute_lib.py", line 3011, in update
    return self._update(var, fn, args, kwargs, group)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/tpu_strategy.py", line 1655, in _update
    fn(value, *distribute_utils.select_replica(i, args),
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py", line 596, in wrapper
    return func(*args, **kwargs)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/distribute/tpu_util.py", line 101, in assign_fn
    handle, ops.convert_to_tensor(value, dtype=var.dtype), name=name)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/profiler/trace.py", line 183, in wrapped
    return func(*args, **kwargs)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/ops.py", line 698, in convert_to_tensor
    return tensor_conversion_registry.convert(
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/tensor_conversion_registry.py", line 234, in convert
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py", line 328, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py", line 267, in constant
    return _constant_impl(value, dtype, shape, name, verify_shape=False,
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py", line 279, in _constant_impl
    return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py", line 289, in _constant_eager_impl
    t = convert_to_eager_tensor(value, ctx, dtype)
  File "/home/ricardo_alonso/venv310/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py", line 102, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
ValueError: No cast function available.
```
