add warning if data parallel model is created for gpus that we dont have
Summary: Don't want to assert since it can be useful to sometimes create models that are not run (for example, unit tests).
Reviewed By: pietern
Differential Revision: D5258905
fbshipit-source-id: f1beee0605bfef235ed0f23f7e78259109720254
diff --git a/caffe2/python/data_parallel_model.py b/caffe2/python/data_parallel_model.py
index 40b1acd..cb659b7 100644
--- a/caffe2/python/data_parallel_model.py
+++ b/caffe2/python/data_parallel_model.py
@@ -74,6 +74,13 @@
'''
if devices is None:
devices = list(range(0, workspace.NumCudaDevices())),
+
+ for gpu in devices:
+ if gpu >= workspace.NumCudaDevices():
+ log.warning("** Only {} GPUs available, GPUs {} requested".format(
+ workspace.NumCudaDevices(), devices))
+ break
+
log.info("Parallelizing model for devices: {}".format(devices))
extra_workers = 8 if rendezvous is not None else 0 # best-guess
num_workers = len(devices) * 4 + extra_workers
