don't import torch in AMD

This commit is contained in:
Somdev Sangwan 2023-05-30 07:24:10 +05:30
parent 251173ffa9
commit a2b7394d8e

8
run.py
View File

@ -1,7 +1,6 @@
#!/usr/bin/env python3
import sys
import time
import torch
import shutil
import core.globals
@ -11,6 +10,11 @@ if not shutil.which('ffmpeg'):
if '--gpu' not in sys.argv:
core.globals.providers = ['CPUExecutionProvider']
if 'ROCMExecutionProvider' not in core.globals.providers:
import torch
if not torch.cuda.is_available():
quit("You are using --gpu flag but CUDA isn't available on your system.")
import glob
import argparse
import multiprocessing as mp
@ -42,8 +46,6 @@ parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_fr
for name, value in vars(parser.parse_args()).items():
args[name] = value
if not torch.cuda.is_available() and args['gpu']:
quit("You are using --gpu flag but CUDA isn't available on your system.")
sep = "/"
if os.name == "nt":