Incorrect input shapes or backend issue with data.triple.Triple and nn.DeepONet?

[voloddia]

Hi,
I run the code below with tensorflow and pytorch backends .

data = dde.data.triple.Triple(
    X_train=(input_trunk, input_branch),
    y_train=output,
    X_test=(input_trunk, input_branch),
    y_test=output,
)

net = dde.nn.DeepONet(
    [input_branch.shape[1], 128],
    [input_trunk.shape[1], 128],
    "relu",
    "Glorot normal",
)

model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["mean l2 relative error"])
losshistory, train_state = model.train(iterations=10000)

This runs without error with tensorflow (although large loss) but has a matrix dimensions mismatch error with pytorch.

input_trunk.shape = (n_samples, d_params)
input_branch.shape = (n_samples, m_steps)
output.shape = (n_samples, m_steps)

Are the shapes correct for data.triple.Triple?