Fix tensorflow import

README.md

@@ -342,7 +342,7 @@ With authors' permission, we listed a set of NNI usage examples and relevant art
 Join IM discussion groups:
 |Gitter||WeChat|
 |----|----|----|
-|![image](https://user-images.githubusercontent.com/39592018/80665738-e0574a80-8acc-11ea-91bc-0836dc4cbf89.png)| OR |![image](https://github.com/JSong-Jia/NNI-user-group/blob/master/user%20group%20code_0512.jpg)|
+|![image](https://user-images.githubusercontent.com/39592018/80665738-e0574a80-8acc-11ea-91bc-0836dc4cbf89.png)| OR |![image](https://github.com/JSong-Jia/NNI-user-group/blob/master/user%20group%20code_0512.png)|
 
 
 ## Related Projects

examples/nas/enas-tf/datasets.py

@@ -2,7 +2,6 @@
 # Licensed under the MIT license.
 
 import tensorflow as tf
-from tensorflow.data import Dataset
 
 def get_dataset():
     (x_train, y_train), (x_valid, y_valid) = tf.keras.datasets.cifar10.load_data()

src/sdk/pynni/nni/nas/tensorflow/enas/trainer.py

@@ -4,7 +4,6 @@
 import logging
 
 import tensorflow as tf
-from tensorflow.data import Dataset
 from tensorflow.keras.optimizers import Adam
 
 from nni.nas.tensorflow.utils import AverageMeterGroup, fill_zero_grads
@@ -39,9 +38,9 @@ def __init__(self, model, loss, metrics, reward_function, optimizer, batch_size,
 
         x, y = dataset_train
         split = int(len(x) * 0.9)
-        self.train_set = Dataset.from_tensor_slices((x[:split], y[:split]))
-        self.valid_set = Dataset.from_tensor_slices((x[split:], y[split:]))
-        self.test_set = Dataset.from_tensor_slices(dataset_valid)
+        self.train_set = tf.data.Dataset.from_tensor_slices((x[:split], y[:split]))
+        self.valid_set = tf.data.Dataset.from_tensor_slices((x[split:], y[split:]))
+        self.test_set = tf.data.Dataset.from_tensor_slices(dataset_valid)
 
         self.mutator = EnasMutator(model)
         self.mutator_optim = Adam(learning_rate=mutator_lr)
@@ -151,9 +150,9 @@ def validate_one_epoch(self, epoch):
 
 
     def _create_train_loader(self):
-        train_set = self.train_set.shuffle(1000000).batch(self.batch_size)
-        test_set = self.test_set.shuffle(1000000).batch(self.batch_size)
+        train_set = self.train_set.shuffle(1000000).repeat().batch(self.batch_size)
+        test_set = self.test_set.shuffle(1000000).repeat().batch(self.batch_size)
         return iter(train_set), iter(test_set)
 
     def _create_validate_loader(self):
-        return iter(self.test_set.shuffle(1000000).batch(self.batch_size))
+        return iter(self.test_set.shuffle(1000000).repeat().batch(self.batch_size))