Update tf_test.py to use jdl.ArrayDataset and 'tensorflow' backend

BirkhoffG · Feb 15, 2024 · dd99bea · dd99bea
1 parent c31ad48
commit dd99bea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/integration_tests/tf_test.py b/integration_tests/tf_test.py
@@ -5,8 +5,8 @@
 
 
 def test_jax():
-    ds = tf.data.Dataset.from_tensor_slices((tf.ones((10, 3)), tf.ones((10, 3))))
-    dl = jdl.DataLoader(ds, 'jax', batch_size=2)
+    ds = jdl.ArrayDataset(np.ones((10, 3)), np.ones((10, 3)))
+    dl = jdl.DataLoader(ds, 'tensorflow', batch_size=2)
     for x, y in dl:
         z = x + y
         assert isinstance(z, np.ndarray)