From 78186fdb1a0169f3d647586d81947fc4ff784c4f Mon Sep 17 00:00:00 2001 From: Quentin Anthony Date: Sat, 23 Apr 2022 09:34:39 -0500 Subject: [PATCH 1/2] Fixed bug. Move dataset to GPU for distributed training --- imagenet/main.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/imagenet/main.py b/imagenet/main.py index 8d41291745..05421d09b6 100644 --- a/imagenet/main.py +++ b/imagenet/main.py @@ -290,7 +290,7 @@ def train(train_loader, model, criterion, optimizer, epoch, args): # measure data loading time data_time.update(time.time() - end) - if args.gpu is not None: + if torch.cuda.is_available(): images = images.cuda(args.gpu, non_blocking=True) if torch.cuda.is_available(): target = target.cuda(args.gpu, non_blocking=True) @@ -334,7 +334,7 @@ def validate(val_loader, model, criterion, args): with torch.no_grad(): end = time.time() for i, (images, target) in enumerate(val_loader): - if args.gpu is not None: + if torch.cuda.is_available(): images = images.cuda(args.gpu, non_blocking=True) if torch.cuda.is_available(): target = target.cuda(args.gpu, non_blocking=True) From 57997c44f36b7ea693a8c2eb5d659faec43f6308 Mon Sep 17 00:00:00 2001 From: Quentin Anthony Date: Tue, 26 Apr 2022 09:49:33 -0500 Subject: [PATCH 2/2] Normalize train and validate GPU movement conditions --- imagenet/main.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/imagenet/main.py b/imagenet/main.py index 05421d09b6..d5731ac1da 100644 --- a/imagenet/main.py +++ b/imagenet/main.py @@ -290,9 +290,9 @@ def train(train_loader, model, criterion, optimizer, epoch, args): # measure data loading time data_time.update(time.time() - end) - if torch.cuda.is_available(): + if args.gpu is not None: images = images.cuda(args.gpu, non_blocking=True) - if torch.cuda.is_available(): + if args.gpu is not None: target = target.cuda(args.gpu, non_blocking=True) # compute output @@ -334,9 +334,9 @@ def validate(val_loader, model, criterion, args): with torch.no_grad(): end = time.time() for i, (images, target) in enumerate(val_loader): - if torch.cuda.is_available(): + if args.gpu is not None: images = images.cuda(args.gpu, non_blocking=True) - if torch.cuda.is_available(): + if args.gpu is not None: target = target.cuda(args.gpu, non_blocking=True) # compute output