def get(self, idx: int) -> Data:
if self.len() == 1:
return copy.copy(self.data)
if not hasattr(self, '_data_list') or self._data_list is None:
self._data_list = self.len() * [None]
elif self._data_list[idx] is not None:
return copy.copy(self._data_list[idx])
data = separate(
cls=self.data.__class__,
batch=self.data,
idx=idx,
slice_dict=self.slices,
decrement=False,
)
self._data_list[idx] = copy.copy(data)
return data
def get_example(self, idx: int) -> BaseData:
r"""Gets the :class:`~torch_geometric.data.Data` or
:class:`~torch_geometric.data.HeteroData` object at index :obj:`idx`.
The :class:`~torch_geometric.data.Batch` object must have been created
via :meth:`from_data_list` in order to be able to reconstruct the
initial object."""
if not hasattr(self, '_slice_dict'):
raise RuntimeError(
("Cannot reconstruct 'Data' object from 'Batch' because "
"'Batch' was not created via 'Batch.from_data_list()'"))
data = separate(
cls=self.__class__.__bases__[-1],
batch=self,
idx=idx,
slice_dict=self._slice_dict,
inc_dict=self._inc_dict,
decrement=True,
)
return data
