@@ -197,10 +197,10 @@ def image_iterator(framework, is_tf_version_2, get_default_mnist_subset, default
197197 return torch .utils .data .DataLoader (dataset = dataset , batch_size = default_batch_size , shuffle = True )
198198
199199 if framework == "mxnet" :
200- from mxnet import gluon
200+ import mxnet
201201
202- dataset = gluon .data .dataset .ArrayDataset (x_train_mnist , y_train_mnist )
203- return gluon .data .DataLoader (dataset , batch_size = 5 , shuffle = True )
202+ dataset = mxnet . gluon .data .dataset .ArrayDataset (x_train_mnist , y_train_mnist )
203+ return mxnet . gluon .data .DataLoader (dataset , batch_size = 5 , shuffle = True )
204204
205205 return None
206206
@@ -323,20 +323,20 @@ def _expected_values():
323323
324324@pytest .fixture (scope = "session" )
325325def get_image_classifier_mx_model ():
326- from mxnet . gluon import nn
326+ import mxnet
327327
328328 # TODO needs to be made parameterizable once Mxnet allows multiple identical models to be created in one session
329329 from_logits = True
330330
331- class Model (nn .Block ):
331+ class Model (mxnet . gluon . nn .Block ):
332332 def __init__ (self , ** kwargs ):
333333 super (Model , self ).__init__ (** kwargs )
334- self .model = nn .Sequential ()
334+ self .model = mxnet . gluon . nn .Sequential ()
335335 self .model .add (
336- nn .Conv2D (channels = 1 , kernel_size = 7 , activation = "relu" ,),
337- nn .MaxPool2D (pool_size = 4 , strides = 4 ),
338- nn .Flatten (),
339- nn .Dense (10 , activation = None ,),
336+ mxnet . gluon . nn .Conv2D (channels = 1 , kernel_size = 7 , activation = "relu" ,),
337+ mxnet . gluon . nn .MaxPool2D (pool_size = 4 , strides = 4 ),
338+ mxnet . gluon . nn .Flatten (),
339+ mxnet . gluon . nn .Dense (10 , activation = None ,),
340340 )
341341
342342 def forward (self , x ):
0 commit comments