fixed test

flyteorg · Oct 29, 2022 · 95cfdaa · 95cfdaa
1 parent 3090609
commit 95cfdaa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/flytekit/unit/extras/tensorflow/test_tensor.py b/tests/flytekit/unit/extras/tensorflow/test_tensor.py
@@ -6,7 +6,7 @@
 
 @task
 def create_rank1_tensor() -> tf.Tensor:
-    return tf.constant([5, 4, 3, 2, 1], dtype=tf.int32)
+    return tf.constant([6, 5, 4, 3, 2, 1], dtype=tf.int32)
 
 
 @task
@@ -33,13 +33,13 @@ def task1(tensor: tf.Tensor) -> tf.Tensor:
 
 @task
 def task2(tensor: tf.Tensor) -> tf.Tensor:
-    return tensor.reshape(tensor, [2, 3])
+    return tf.reshape(tensor, [2, 3])
 
 
 @workflow
 def wf():
     task1(tensor=create_rank1_tensor())
-    # task2(tensor=create_rank1_tensor())
+    task2(tensor=create_rank1_tensor())
 
 
 @workflow