Browse Source
Fix floating point in number of workers `nw` (#6701)
Integer division by a float yields a (rounded) float. This causes the dataloader to crash when creating a range.modifyDataloader
Samuel Yvon
GitHub
2 years ago
1 changed files with 1 additions and 1 deletions
+ 1
- 1
utils/datasets.py
View File
| batch_size = min(batch_size, len(dataset)) | batch_size = min(batch_size, len(dataset)) | ||||
| nd = torch.cuda.device_count() # number of CUDA devices | nd = torch.cuda.device_count() # number of CUDA devices | ||||
| nw = min([os.cpu_count() // max(nd / 2, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers | |||||
| nw = min([os.cpu_count() // max(nd // 2, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers | |||||
| sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) | sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) | ||||
| loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates | loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates | ||||
| return loader(dataset, | return loader(dataset, |
Loading…