Dataset Docs updates (#13015)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2024-05-22 12:17:32 +02:00 committed by GitHub
parent ade1743fa5
commit 2dc0a62bb1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 6 additions and 17 deletions

View file

@ -128,7 +128,7 @@ def build_dataloader(dataset, batch, workers, shuffle=True, rank=-1):
"""Return an InfiniteDataLoader or DataLoader for training or validation set."""
batch = min(batch, len(dataset))
nd = torch.cuda.device_count() # number of CUDA devices
nw = min([os.cpu_count() // max(nd, 1), workers]) # number of workers
nw = min(os.cpu_count() // max(nd, 1), workers) # number of workers
sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
generator = torch.Generator()
generator.manual_seed(6148914691236517205 + RANK)