to_dense_batch raises IndexError: index is out of bounds for dimension with size 0 on GPU

[yoyoong]

I encountered a runtime error when calling to_dense_batch on GPU, while the same code works correctly on CPU:

File "/home/github/RetroBridge/src/frameworks/markov_bridge.py", line 177, in training_step
reactants, product, pred, node_mask, noisy_data, _ = self.process_and_forward(data)
File "/home/github/RetroBridge/src/frameworks/markov_bridge.py", line 144, in process_and_forward
reactants, r_node_mask = utils.to_dense(data.x, data.edge_index, data.edge_attr, data.batch)
File "/home/github/RetroBridge/src/data/utils.py", line 81, in to_dense
X, node_mask = to_dense_batch(x=x, batch=batch)
File "/home/anaconda3/envs/retrobridge/lib/python3.9/site-packages/torch_geometric/experimental.py", line 117, in wrapper
return func(*args, **kwargs)
File "/home/anaconda3/envs/retrobridge/lib/python3.9/site-packages/torch_geometric/utils/_to_dense_batch.py", line 128, in to_dense_batch
out[idx] = x
IndexError: index is out of bounds for dimension with size 0

I don't change anything of the code.
I would really appreciate it if someone could help take a look at this issue.