address comments

intel-analytics · Feb 9, 2018 · e6658a2 · e6658a2
1 parent b9c3e5a
commit e6658a2
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/pyspark/bigdl/models/utils/model_broadcast.py b/pyspark/bigdl/models/utils/model_broadcast.py
@@ -21,13 +21,11 @@
 
 from pyspark.cloudpickle import print_exec
 from pyspark.broadcast import Broadcast
+from pyspark.broadcast import _from_id
 from bigdl.nn.layer import Model
 
 def _from_id_and_type(bid, bigdl_type):
-    from pyspark.broadcast import _broadcastRegistry
-    if bid not in _broadcastRegistry:
-        raise Exception("Broadcast variable '%s' not loaded!" % bid)
-    result = _broadcastRegistry[bid]
+    result = _from_id(bid)
     return ModelBroadcast(path=result._path, bigdl_type=bigdl_type)
 
 def broadcastModel(sc, layer):