Repository: pyg-team/pytorch_geometric

Issue:
Incorrect batching with graphs converted by `ToSparseTensor()`
### 🐛 Describe the bug

I've been debugging my code and have run into an issue with my custom dataset - at present, I have it return a `HeteroData` object that has been converted with the `ToSparseTensor()` transform to construct the sparse form (`adj_t`). I then feed this dataset into the torch_geometric custom dataloader (`torch_geometric.loader.DataLoader`).

However, when I get batches from this dataloader, the `adj_t` of the batched graph seems wildly incorrect (that is, it appears to properly extend the sparse adjacency matrix along just one dimension, but not both, leading to incorrect indices). 

Minimally reproducible code snippet:

```
import torch
import torch_geometric

rng = torch.Generator().manual_seed(9)
bsize = 2

h_dataset = torch_geometric.datasets.FakeHeteroDataset(num_graphs=128, 
                                                  num_node_types=4, 
                                                  num_edge_types=16, 
                                                  avg_num_nodes=1024, 
                                                  num_classes=2, 
                                                  transform=torch_geometric.transforms.ToSparseTensor())


train_ds, test_ds = torch.utils.data.random_split(h_dataset, [64, 64], rng)

for i in range(bsize):
    print(train_ds[i])

train_dl = torch_geometric.loader.DataLoader(train_ds, shuffle=False, batch_size=2, drop_last=False)
batched_data = next(iter(train_dl))
print(batched_data)
```

Specifically, you'll see that the printing of the first two graphs in the train_ds (before batching) produces values around 1000 for the number of each node type and the adjacency matrices match appropriately. 

For example, from my environment, the edge data contains: `(v2, e0, v2)={ adj_t=[1237, 1237] }` and `(v2, e0, v2)={ adj_t=[1005, 1005] }` for the two graphs.

After batching, the edge data contains: `(v2, e0, v2)={ adj_t=[2242, 1237] }` - which appears as if the batching worked along the first dimension but not the second.

### Environment

* PyG version: 2.4.0
* PyTorch version: 2.1.0
* OS: Mac OSX 
* Python version: 3.10.13
* CUDA/cuDNN version: N/A
* How you installed PyTorch and PyG (`conda`, `pip`, source): pip
* Any other relevant information (*e.g.*, version of `torch-scatter`): N/A



Discussion:
~~While trying to figure out what was happening here, I think it has to do somehow with the collate function within the dataloader - specifically in the Collater class here (https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/loader/dataloader.html#DataLoader) which just calls https://github.com/pytorch/pytorch/blob/main/torch/utils/data/_utils/collate.py - but this doesn't seem to behave appropriately with the sparse tensor/adjacency format that is used here.~~

Never mind - I misread the code fundamentally. `Collater` (https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/loader/dataloader.html#Collater) calls 

`Batch.from_data_list()` (https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/data/batch.html#Batch.from_data_list) which then calls 

`torch_geometric.data.collate()` here (https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/data/collate.py). 

I haven't dug deeply enough into `collate` to figure out where things are going wrong, but I am trying to write my own custom `collate` function to pass into a vanilla `DataLoader` to try to workaround this for now. I will report back if I find the specific issue. 
I think I've localized the issue - `__cat_dim__` in HeteroData is defined to return (0,1) if the data is sparse and is the adjacency matrix; however, this check does not work if torch_sparse is not installed as it specifically checks for a `SparseTensor` from torch_sparse rather than any instance of a sparse tensor:

https://github.com/pyg-team/pytorch_geometric/blob/999af2312efcdd7bd35e40c9b21c4a6a45464ac5/torch_geometric/data/hetero_data.py#L335C17-L335C17

Resolving patch:
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -28,6 +28,7 @@
 from torch_geometric.utils import (
     bipartite_subgraph,
     contains_isolated_nodes,
+    is_sparse,
     is_undirected,
     mask_select,
 )
@@ -332,7 +333,7 @@ def update(self, data: 'HeteroData') -> 'HeteroData':
     def __cat_dim__(self, key: str, value: Any,
                     store: Optional[NodeOrEdgeStorage] = None, *args,
                     **kwargs) -> Any:
-        if isinstance(value, SparseTensor) and 'adj' in key:
+        if is_sparse(value) and 'adj' in key:
             return (0, 1)
         elif isinstance(store, EdgeStorage) and 'index' in key:
             return -1


Test patch: