ValueError: MessagePassing.propagate only supports torch.LongTensor of shape [2, num_messages] or torch_sparse.SparseTensor for argument edge_index.

[vrjmath]

====epoch 1
Iteration: 0%| | 0/38 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/usr/scratch/vshitole6/Summer/pretraingnns/pretraingnns/chem/finetune.py", line 232, in
main()
File "/usr/scratch/vshitole6/Summer/pretraingnns/pretraingnns/chem/finetune.py", line 204, in main
train(args, model, device, train_loader, optimizer)
File "/usr/scratch/vshitole6/Summer/pretraingnns/pretraingnns/chem/finetune.py", line 32, in train
pred = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/Summer/pretraingnns/pretraingnns/chem/model.py", line 367, in forward
node_representation = self.gnn(x, edge_index, edge_attr)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/Summer/pretraingnns/pretraingnns/chem/model.py", line 268, in forward
h = self.gnns[layer](h_list[layer], edge_index, edge_attr)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/Summer/pretraingnns/pretraingnns/chem/model.py", line 49, in forward
return self.propagate(self.aggr, edge_index, x=x, edge_attr=edge_embeddings)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch_geometric/nn/conv/message_passing.py", line 216, in propagate
size = self.check_input(edge_index, size)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/scratch/vshitole6/anaconda3/lib/python3.11/site-packages/torch_geometric/nn/conv/message_passing.py", line 110, in check_input
raise ValueError(
ValueError: MessagePassing.propagate only supports torch.LongTensor of shape [2, num_messages] or torch_sparse.SparseTensor for argument edge_index.