fixup! Add support for TensorFlow 2.16

nengo · Apr 4, 2024 · 0034bda · 0034bda
1 parent ab071f2
commit 0034bda
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -40,7 +40,7 @@ jobs:
         include:
           - script: test
             coverage-name: gpu
-            tf-version: tensorflow==2.10
+            tf-version: tensorflow[and-cuda]
             runs-on: [self-hosted, gpu]
           - script: test
             python-version: "3.11"
@@ -50,11 +50,11 @@ jobs:
             python-version: "3.8"
             coverage-name: oldest
           - script: docs
-            tf-version: tensorflow==2.10
+            tf-version: tensorflow[and-cuda]==2.16.1
             python-version: "3.9"
             runs-on: [self-hosted, gpu]
           - script: examples
-            tf-version: tensorflow==2.10
+            tf-version: tensorflow[and-cuda]==2.16.1
             runs-on: [self-hosted, gpu]
       fail-fast: false
     env:

diff --git a/keras_lmu/tests/test_benchmarks.py b/keras_lmu/tests/test_benchmarks.py
@@ -39,13 +39,13 @@ def on_predict_batch_end(self, batch, logs=None):
 @pytest.mark.skipif(not tf_gpu_installed, reason="Very slow on CPU")
 @pytest.mark.parametrize(
     "mode, min_time, max_time",
-    [("rnn", 0.1, 0.2), ("fft", 0.05, 0.15), ("raw", 0.05, 0.15)],
+    [("rnn", 0.01, 0.1), ("fft", 0.01, 0.1), ("raw", 0.01, 0.1)],
 )
 def test_performance(mode, min_time, max_time):
-    # performance is based on Azure NC6 VM
-    # CPU: Intel Xeon E5-2690 v3 @ 2.60Ghz
-    # GPU: Nvidia Tesla K80
-    # TensorFlow version: 2.6.0
+    # performance is based on
+    # CPU: AMD Ryzen 9 5950X
+    # GPU: Nvidia RTX 3060
+    # TensorFlow version: 2.10.0
 
     dims = 32
     seq_len = 512