[Embedding] Fix the feature filter UT in feature_column_v2_test. (#19)

DeepRec-AI · Dec 29, 2021 · 5cffe8b · 5cffe8b
1 parent 7dac910
commit 5cffe8b
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/tensorflow/python/feature_column/feature_column_v2_test.py b/tensorflow/python/feature_column/feature_column_v2_test.py
@@ -7577,13 +7577,17 @@ def testEmbeddingVariableForFeatureFilter(self):
       sess.run(ops.get_collection(ops.GraphKeys.EV_INIT_SLOT_OPS))
       sess.run([init])
       emb1, top, l = sess.run([emb, train_op, loss])
+      for val1 in emb1.tolist():
+        for val in val1:
+          self.assertEqual(val, 1.0)
       emb1, top, l = sess.run([emb, train_op, loss])
-      emb1, top, l = sess.run([emb, train_op, loss])
-      for val in emb1.tolist()[0]:
-        self.assertEqual(val, 1.0)
-      emb1, top, l = sess.run([emb, train_op, loss])
-      for val in emb1.tolist()[0]:
-        self.assertNotEqual(val, 1.0)
+      for index, val1 in enumerate(emb1.tolist()):
+        if index < 7:
+          for val in val1:
+            self.assertNotEqual(val, 1.0)
+        else:
+          for val in val1:
+            self.assertEqual(val, 1.0)
 
   @test_util.run_deprecated_v1
   def testEmbeddingVariableForAdaptiveEmbedding(self):