recover previous behaviour

n3fit/src/n3fit/backends/keras_backend/callbacks.py

@@ -75,9 +75,15 @@ def __init__(self, stopping_object, log_freq=100):
         super().__init__()
         self.log_freq = log_freq
         self.stopping_object = stopping_object
+        self._current_loss = None
+
+    def on_epoch_begin(self, epoch, logs=None):
+        # TODO This is an unnecessary performance hit, just for testing
+        self._current_loss = self.model.compute_losses()
 
     def on_epoch_end(self, epoch, logs=None):
         """ Function to be called at the end of every epoch """
+        logs = self._current_loss
         print_stats = ((epoch + 1) % self.log_freq) == 0
         # Note that the input logs correspond to the fit before the weights are updated
         self.stopping_object.monitor_chi2(logs, epoch, print_stats=print_stats)

n3fit/src/n3fit/backends/keras_backend/constraints.py

@@ -18,7 +18,6 @@ def __init__(self, min_value, max_value, **kwargs):
             min_value=min_value, max_value=max_value, **kwargs
         )
 
-    @tf.function
     def __call__(self, w):
         norms = K.sum(w, axis=self.axis, keepdims=True)
         desired = (