incompatible shapes during computing loss with both ragged and padded inputs

[ekoag]

Hello,
I am attempting to fit the the Schnet model for a node prediction task but I am experiencing some difficulty. I have used GraphDict and GraphMemoryList to build a dataset of graphs. Each graph is a protein molecule. Here are example dimensions of a single graph:
node_attributes (1337, 12)
edge_attributes (21392, 1)
edge_indices (21392, 2)
node_labels (1337, 1)

I have tried both ragged and padded inputs but the model fails to compute loss once the batch size is greater than 1.
Here is my model configuration, passed into Schnet's make_model:
model_config = {
"inputs" : [
{"shape": (None, 12), "name": "node_attributes", "dtype": "float32"},
{"shape": (None, 1), "name": "edge_attributes", "dtype": "float32"},
{"shape": (None, 2), "name": "edge_indices", "dtype": "int64"},
{"shape": (), "name": "total_nodes", "dtype": "float32"},
{"shape": (), "name": "total_edges", "dtype": "float32"}],
"depth": 4,
"input_node_embedding": {"input_dim": 12 , "output_dim": 256},
"make_distance": False,
"expand_distance": True,
"input_tensor_type": "padded",
"output_embedding": "node",}

When training Schnet, I encounter the following error:
1/2 ━━━━━━━━━━━━━━━━━━━━ 40s 41s/step - loss: 2.4767 - mean_absolute_error: 2.47672024-12-16 14:17:51.739649: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Incompatible shapes: [18,14565,1] vs. [18,4815,1]
Traceback (most recent call last):
File "/Users/ekoag/kgcnn/test.py", line 110, in
model.fit(x_train,labels,epochs=epochs,callbacks=callbacks)
File "/Users/ekoag/Library/Python/3.9/lib/python/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/Users/ekoag/Library/Python/3.9/lib/python/site-packages/tensorflow/python/framework/ops.py", line 5983, in raise_from_not_ok_status
raise core._status_to_exception(e) from None # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node _wrapped__Sub_device/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [18,14565,1] vs. [18,4815,1] [Op:Sub] name: