github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 21:12:07 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix for colab

Browse source
This commit is contained in:
iperov 2021-04-22 21:12:42 +04:00
parent 78f12de819
commit 23130cd56a
1 changed files with 8 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

15
core/leras/device.py
View file

@ -92,12 +92,13 @@ class Devices(object):
@staticmethod @staticmethod
def _get_tf_devices_proc(q : multiprocessing.Queue): def _get_tf_devices_proc(q : multiprocessing.Queue):
compute_cache_path = Path(os.environ['APPDATA']) / 'NVIDIA' / ('ComputeCache_ALL') if sys.platform[0:3] == 'win':
os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path) compute_cache_path = Path(os.environ['APPDATA']) / 'NVIDIA' / ('ComputeCache_ALL')
if not compute_cache_path.exists(): os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path)
io.log_info("Caching GPU kernels...") if not compute_cache_path.exists():
compute_cache_path.mkdir(parents=True, exist_ok=True) io.log_info("Caching GPU kernels...")
compute_cache_path.mkdir(parents=True, exist_ok=True)
import tensorflow import tensorflow
tf_version = tensorflow.version.VERSION tf_version = tensorflow.version.VERSION
@ -156,7 +157,7 @@ class Devices(object):
def initialize_main_env(): def initialize_main_env():
if int(os.environ.get("NN_DEVICES_INITIALIZED", 0)) != 0: if int(os.environ.get("NN_DEVICES_INITIALIZED", 0)) != 0:
return return
if 'CUDA_VISIBLE_DEVICES' in os.environ.keys(): if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
os.environ.pop('CUDA_VISIBLE_DEVICES') os.environ.pop('CUDA_VISIBLE_DEVICES')