Wraps the :attr:`worker_init_fn` argument of the DataLoader to set the number of DGL OMP threads to 1 for PyTorch DataLoader workers.
| 741 | |
| 742 | |
| 743 | class WorkerInitWrapper(object): |
| 744 | """Wraps the :attr:`worker_init_fn` argument of the DataLoader to set the number of DGL |
| 745 | OMP threads to 1 for PyTorch DataLoader workers. |
| 746 | """ |
| 747 | |
| 748 | def __init__(self, func): |
| 749 | self.func = func |
| 750 | |
| 751 | def __call__(self, worker_id): |
| 752 | set_num_threads(1) |
| 753 | if self.func is not None: |
| 754 | self.func(worker_id) |
| 755 | |
| 756 | |
| 757 | def create_tensorized_dataset( |