search

zudi-lin / pytorch_connectomics

PyTorch Connectomics: segmentation toolbox for EM connectomics
http://connectomics.readthedocs.io/
MIT License
169 stars 77 forks source link

add cpu option for inference #121

Closed zengyuy closed 1 year ago

zengyuy commented 1 year ago

add --use_cpu option for model inference

zudi-lin commented 1 year ago

Hi @zengyuy, could you update the connectomics/utils/system.py file without the use_cpu option? I think 

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

is sufficient as we want to use GPU by default if it is available. Other changes look good to me. Please update your PR and I will merge the code. Thanks!