Add mask device, fix test

Reviewed By: azzolini Differential Revision: D5930258 fbshipit-source-id: 16fdc2aeba7d95e815e55ca495118a5129495bb0
achraf-mer · Sep 28, 2017 · b9009df · b9009df
1 parent 642dea4
commit b9009df
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/caffe2/python/core.py b/caffe2/python/core.py
@@ -78,10 +78,12 @@ def IsOperatorWithEngine(op_type, engine):
     return (op_type + "_ENGINE_" + engine in _REGISTERED_OPERATORS)
 
 
-def DeviceOption(device_type, cuda_gpu_id=0, random_seed=None):
+def DeviceOption(device_type, cuda_gpu_id=0, random_seed=None, node_name=None):
     option = caffe2_pb2.DeviceOption()
     option.device_type = device_type
     option.cuda_gpu_id = cuda_gpu_id
+    if node_name is not None:
+        option.node_name = node_name
     if random_seed is not None:
         option.random_seed = random_seed
     return option