Merge pull request #100 from hzdr-MedImaging/fix-cuda-init

remove obsolete cuda init which is performed in next if anyway.
wolny · Dec 4, 2023 · 725ca3a · 725ca3a
2 parents a296dc9 + e39b51c
commit 725ca3a
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 2 deletions.
diff --git a/pytorch3dunet/predict.py b/pytorch3dunet/predict.py
@@ -38,7 +38,6 @@ def main():
     if torch.cuda.device_count() > 1 and not config['device'] == 'cpu':
         model = nn.DataParallel(model)
         logger.info(f'Using {torch.cuda.device_count()} GPUs for prediction')
-        model = model.cuda()
     if torch.cuda.is_available() and not config['device'] == 'cpu':
         model = model.cuda()
 

diff --git a/pytorch3dunet/unet3d/trainer.py b/pytorch3dunet/unet3d/trainer.py
@@ -23,7 +23,6 @@ def create_trainer(config):
     if torch.cuda.device_count() > 1 and not config['device'] == 'cpu':
         model = nn.DataParallel(model)
         logger.info(f'Using {torch.cuda.device_count()} GPUs for prediction')
-        model = model.cuda()
     if torch.cuda.is_available() and not config['device'] == 'cpu':
         model = model.cuda()