diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py index 8b461945466b7bacd300618cb2ddc420eccdc367..502676b43b38dea7bb7cfac4a2b3c08999cbc62a 100644 --- a/pysegcnn/core/trainer.py +++ b/pysegcnn/core/trainer.py @@ -2106,7 +2106,7 @@ class NetworkInference(BaseConfig): # the device to compute on, use gpu if available self.device = torch.device( - 'cuda:' if torch.cuda.is_available() else 'cpu') + 'cuda:0' if torch.cuda.is_available() else 'cpu') # the output paths for the different graphics self.base_path = pathlib.Path(HERE)