diff --git a/keras_lmu/layers.py b/keras_lmu/layers.py
index 4680c524..e5efd983 100644
--- a/keras_lmu/layers.py
+++ b/keras_lmu/layers.py
@@ -1,6 +1,4 @@
-"""
-Core classes for the KerasLMU package.
-"""
+"""Core classes for the KerasLMU package."""
 
 import numpy as np
 import tensorflow as tf
@@ -180,8 +178,8 @@ def theta(self):
         """
         Value of the ``theta`` parameter.
 
-        If ``trainable_theta=True`` this returns the trained value, not the initial
-        value passed in to the constructor.
+        If ``trainable_theta=True`` this returns the trained value, not the
+        initial value passed in to the constructor.
         """
         if self.built:
             return 1 / tf.keras.backend.get_value(self.theta_inv)
@@ -583,8 +581,8 @@ def theta(self):
         """
         Value of the ``theta`` parameter.
 
-        If ``trainable_theta=True`` this returns the trained value, not the initial
-        value passed in to the constructor.
+        If ``trainable_theta=True`` this returns the trained value, not the
+        initial value passed in to the constructor.
         """
 
         if self.built:
diff --git a/keras_lmu/tests/test_benchmarks.py b/keras_lmu/tests/test_benchmarks.py
index d58de82e..dffb1832 100644
--- a/keras_lmu/tests/test_benchmarks.py
+++ b/keras_lmu/tests/test_benchmarks.py
@@ -38,7 +38,7 @@ 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.1, 0.2), ("raw", 0.05, 0.15)],
+    [("rnn", 0.1, 0.2), ("fft", 0.05, 0.15), ("raw", 0.05, 0.15)],
 )
 def test_performance(mode, min_time, max_time):
     # performance is based on Azure NC6 VM
diff --git a/keras_lmu/tests/test_layers.py b/keras_lmu/tests/test_layers.py
index 509afed4..7590a21e 100644
--- a/keras_lmu/tests/test_layers.py
+++ b/keras_lmu/tests/test_layers.py
@@ -82,7 +82,9 @@ def test_layer_vs_cell(rng, has_input_kernel, feedforward, discretizer):
         kernel_initializer="glorot_uniform" if has_input_kernel else None,
         memory_to_memory=not feedforward,
     )
-    hidden_cell = lambda: tf.keras.layers.SimpleRNNCell(units=64)
+
+    def hidden_cell():
+        return tf.keras.layers.SimpleRNNCell(units=64)
 
     inp = rng.uniform(-1, 1, size=(2, n_steps, input_d))