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

remove obsolete cuda init which is performed in next if anyway.

Author: wolny

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()
 

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()