Add exit when GPU not found

2020-04-20 19:27:57 -04:00 · 2020-04-20 19:27:57 -04:00 · 5adf6cbac6
commit 5adf6cbac6
parent facfc13132
1 changed files with 3 additions and 0 deletions
--- a/map2map/train.py
+++ b/map2map/train.py
@ -35,6 +35,9 @@ def node_worker(args):

    node = int(os.environ['SLURM_NODEID'])

+    if args.gpus_per_node < 1:
+        raise RuntimeError('GPU not found on node {}'.format(node))
+
    spawn(gpu_worker, args=(node, args), nprocs=args.gpus_per_node)