Can't batch along new dimension for attribute 'x'

[Bobbubble]

Hi,

I'm creating a new HeteroData, and I'm following the tutorial to overwrite the cat_dim() method so that I can batch along a new dimension. In the tutorial, I've managed to successfully batch along a new dimension on Data's 'x' attribute to batch along a new dimension. But when I extend the tutorial code to HeteroData, it doesn't seem to batch along a new dimension on HeteroData's 'x' attribute properly. Here's my implementation:

import torch
from torch_geometric.data import HeteroData
from torch_geometric.loader import DataLoader

class MyHeteroData(HeteroData):
    def __cat_dim__(self, key, value, *args, **kwargs):
        if key == 'x':
            return None
        return super().__cat_dim__(key, value, *args, **kwargs)

data = MyHeteroData()

data['paper'].x = torch.randn(5, 32)
data['author'].x = torch.randn(3, 32)
data['paper', 'cites', 'paper'].edge_index = torch.tensor([
    [0, 1, 2, 3],
    [1, 2, 3, 4],
], dtype=torch.long)

data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)

And I also get the following errors:

  File "C:\Users\User\Desktop\111\batchtest_1.py", line 22, in <module>
    batch = next(iter(loader))
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch\utils\data\dataloader.py", line 530, in __next__
    data = self._next_data()
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch\utils\data\dataloader.py", line 570, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch\utils\data\_utils\fetch.py", line 52, in fetch
    return self.collate_fn(data)
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch_geometric\loader\dataloader.py", line 20, in __call__
    self.exclude_keys)
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch_geometric\data\batch.py", line 82, in from_data_list
    exclude_keys=exclude_keys,
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch_geometric\data\collate.py", line 107, in collate
    repeats = [store.num_nodes for store in stores]
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch_geometric\data\collate.py", line 107, in <listcomp>
    repeats = [store.num_nodes for store in stores]
  File "E:\anaconda3\envs\difusconew\lib\site-packages\torch_geometric\data\storage.py", line 275, in num_nodes
    return value.shape[cat_dim]
TypeError: tuple indices must be integers or slices, not NoneType

Is there a way to help me batch along new dimension for HeteroData' attribute 'x' ?