From 5adf6cbac641ebe4a7c928216dcec95f6406883e Mon Sep 17 00:00:00 2001
From: Yin Li <eelregit@gmail.com>
Date: Mon, 20 Apr 2020 19:27:57 -0400
Subject: [PATCH] Add exit when GPU not found

---
 map2map/train.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/map2map/train.py b/map2map/train.py
index dfcfe88..046e648 100644
--- 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)