Chronos: fix ZooTestCase in tensorflow unit tests (#5800)

intel-analytics · Sep 19, 2022 · de9fa8c · de9fa8c
1 parent 0e31b5a
commit de9fa8c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/python/chronos/test/bigdl/chronos/model/tf2/test_Seq2Seq_keras.py b/python/chronos/test/bigdl/chronos/model/tf2/test_Seq2Seq_keras.py
@@ -19,7 +19,7 @@
 
 import pytest
 
-from bigdl.orca.test_zoo_utils import ZooTestCase
+from unittest import TestCase
 from bigdl.chronos.model.tf2.Seq2Seq_keras import LSTMSeq2Seq, model_creator
 import numpy as np
 
@@ -42,7 +42,7 @@ def get_x_y(num_samples):
 
 
 @pytest.mark.skipif(tf.__version__ < '2.0.0', reason="Run only when tf>2.0.0.")
-class TestSeq2Seq(ZooTestCase):
+class TestSeq2Seq(TestCase):
 
     train_data, test_data = create_data()
     model = model_creator(config={

diff --git a/python/chronos/test/bigdl/chronos/model/tf2/test_tcn_keras.py b/python/chronos/test/bigdl/chronos/model/tf2/test_tcn_keras.py
@@ -19,7 +19,7 @@
 import pytest
 import tensorflow as tf
 
-from bigdl.orca.test_zoo_utils import ZooTestCase
+from unittest import TestCase
 from bigdl.chronos.model.tf2.TCN_keras import model_creator, TemporalConvNet, TemporalBlock
 
 
@@ -40,7 +40,7 @@ def get_x_y(num_samples):
     return train_data, test_data
 
 @pytest.mark.skipif(tf.__version__ < '2.0.0', reason="Run only when tf>2.0.0.")
-class TestTcnKeras(ZooTestCase):
+class TestTcnKeras(TestCase):
 
     train_data, test_data = create_data()
     model = model_creator(config={