fix dataset tests for JAX

bayesflow-org · Jun 19, 2024 · f388b74 · f388b74
1 parent cf8d420
commit f388b74
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_datasets/conftest.py b/tests/test_datasets/conftest.py
@@ -20,11 +20,11 @@ class Model(keras.Model):
         def call(self, *args, **kwargs):
             pass
 
-        def compute_loss(self, **kwargs):
+        def compute_loss(self, *args, **kwargs):
             return keras.ops.zeros(())
 
     model = Model()
-    model.compile(optimizer=None)
+    model.compile()
 
     return model