CalvinBackup/SimSwapPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update dataloader

Browse Source
This commit is contained in:
chenxuanhong
2022-02-15 02:08:57 +08:00
parent 4a6197a685
commit 9c1e6bb77d
3 changed files with 11 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data_tools/data_loader_VGGFace2HQ_multigpu.py
+2 -2
View File
@@ -5,7 +5,7 @@
# Created Date: Sunday February 6th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Tuesday, 15th February 2022 1:35:41 am
# Last Modified: Tuesday, 15th February 2022 1:50:19 am
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
@@ -56,7 +56,7 @@ class InfiniteSampler(torch.utils.data.Sampler):
class data_prefetcher():
def __init__(self, loader, cur_gpu):
torch.cuda.set_device(cur_gpu)
torch.cuda.set_device(cur_gpu) # must add this line to avoid excessive use of GPU 0 by the prefetcher
self.loader = loader
self.dataiter = iter(loader)
self.stream = torch.cuda.Stream(device=cur_gpu)