import os
os.environ["KERAS_BACKEND"] = "torch"
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import numpy as np
import keras

inputs = keras.Input(shape=(10, 16))
x = keras.layers.MultiHeadAttention(num_heads=2, key_dim=16)(inputs, inputs)
outputs = keras.layers.Dense(8)(x)
model = keras.Model(inputs=inputs, outputs=outputs)

model.compile(optimizer="adam", loss="mse")

X_dummy = np.random.rand(5, 10, 16).astype(np.float32)
y_dummy = np.random.rand(5, 10, 8).astype(np.float32)

model.fit(X_dummy, y_dummy, epochs=1, verbose=1)

spec = [keras.InputSpec(name="input", shape=(1, 10, 16), dtype="float32")]
model.export("test.onnx", format="onnx", verbose=False, input_signature=spec)