Commit 41c8446c authored by Yuxin Wu's avatar Yuxin Wu

add missing file

parent eee59532
......@@ -136,9 +136,9 @@ class Trainer(object):
except (KeyboardInterrupt, Exception):
raise
finally:
self.coord.request_stop()
# Do I need to run queue.close?
callbacks.after_train()
self.coord.request_stop()
self.summary_writer.close()
self.sess.close()
......
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: gpu.py
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>
from .utils import change_env
__all__ = ['change_gpu', 'get_nr_gpu', 'get_gpus']
def change_gpu(val):
return change_env('CUDA_VISIBLE_DEVICES', str(val))
def get_nr_gpu():
env = os.environ.get('CUDA_VISIBLE_DEVICES', None)
assert env is not None # TODO
return len(env.split(','))
def get_gpus():
""" return a list of GPU physical id"""
env = os.environ.get('CUDA_VISIBLE_DEVICES', None)
assert env is not None # TODO
return map(int, env.strip().split(','))
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment