def torch_dtype(self):
        """Return the data type of the ID map."""
        # [TODO][Rui] Use torch instead of numpy.
        return torch.int32 if self.dtype == np.int32 else torch.int64
Return the data type of the ID map.
python/dgl/distributed/id_map.py