Fix basic_layers.py import

apache · May 21, 2019 · 3615916 · 3615916
1 parent 8848226
commit 3615916
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/mxnet/gluon/contrib/nn/basic_layers.py b/python/mxnet/gluon/contrib/nn/basic_layers.py
@@ -24,7 +24,7 @@
            'PixelShuffle3D']
 
 import warnings
-from .... import nd, test_utils
+from .... import nd, context
 from ...block import HybridBlock, Block
 from ...nn import Sequential, HybridSequential, BatchNorm
 
@@ -233,7 +233,7 @@ def _get_num_devices(self):
         warnings.warn("Caution using SyncBatchNorm: "
                       "if not using all the GPUs, please mannually set num_devices",
                       UserWarning)
-        num_devices = mx.context.num_gpus()
+        num_devices = context.num_gpus()
         num_devices = num_devices if num_devices > 0 else 1
         return num_devices