kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

add cuda check

Browse Source
This commit is contained in:
Qing 2023-01-14 22:00:34 +08:00
parent 719c7278b4
commit 7d00fc8ceb
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lama_cleaner/parse_args.py
View File

@ -75,6 +75,11 @@ def parse_args():
args = parser.parse_args() args = parser.parse_args()
if args.device == "cuda":
import torch
if torch.cuda.is_available() is False:
parser.error("torch.cuda.is_available() is False, please use --device cpu or check your pytorch installation")
if args.model_dir is not None: if args.model_dir is not None:
if os.path.isfile(args.model_dir): if os.path.isfile(args.model_dir):
parser.error(f"invalid --model-dir: {args.model_dir} is a file") parser.error(f"invalid --model-dir: {args.model_dir} is a file")