From 1aba4280d82a9fee8a9f649412afc898bcd7471e Mon Sep 17 00:00:00 2001
From: Byung-Gon Chun <bgchun@fb.com>
Date: Fri, 2 Dec 2016 10:58:39 -0800
Subject: [PATCH] Make xray net_type configurable

Summary: Make xray net_type configub a command line argument

Differential Revision: D4262076

fbshipit-source-id: e2ecb9cd5bee5d6aaebe0ea8d2d4d9b378058cba
---
 caffe2/python/data_parallel_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/caffe2/python/data_parallel_model.py b/caffe2/python/data_parallel_model.py
index 37a5a538f3d..515f266d88a 100644
--- a/caffe2/python/data_parallel_model.py
+++ b/caffe2/python/data_parallel_model.py
@@ -21,6 +21,7 @@ def Parallelize_GPU(
     param_update_builder_fun,
     devices=range(0, workspace.NumCudaDevices()),
     rendezvous=None,
+    net_type='dag',
 ):
     '''
     Function to create a model that can run on many GPUs.
@@ -48,12 +49,13 @@ def Parallelize_GPU(
       rendezvous:       used for rendezvous in distributed computation, if None
                         then only one node is used. To create rendezvous,
                         use <TBD>.
+      net_type:         Network type
 
     '''
     log.info("Parallelizing model for devices: {}".format(devices))
     extra_workers = 8 if rendezvous is not None else 0  # best-guess
     model_helper_obj.net.Proto().num_workers = len(devices) * 2 + extra_workers
-    model_helper_obj.net.Proto().type = 'dag'
+    model_helper_obj.net.Proto().type = net_type
 
     # Store some information in the model -- a bit ugly
     model_helper_obj._devices = devices