From a516f418b966758135068c5bd34706f974033ae2 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 20 Apr 2026 15:40:55 +0200 Subject: [PATCH] Fix a high priority bug where reconstructions will crash upon saving, and add test coverage to prevent a repeat --- src/cdtools/models/base.py | 5 ++++- tests/models/test_fancy_ptycho.py | 3 +++ tests/models/test_simple_ptycho.py | 3 +++ 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 554bd39..5b990a0 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -192,7 +192,10 @@ class CDIModel(t.nn.Module): 'loss_history': np.array(self.loss_history), 'epoch': self.epoch, 'training_history': self.training_history, - 'loss_function': self.loss.__name__, + 'loss_function': ( + getattr(self.loss, '__name__', None) or + getattr(self.loss.func, '__name__', str(self.loss)) + ), } diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index 93b3280..23e2334 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -123,6 +123,9 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): time.sleep(3) plt.close('all') + # Simply test that this does not fail + results = model.save_results(dataset) + # If this fails, the reconstruction has gotten worse assert model.loss_history[-1] < 0.38 diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index bcc5e46..abace7e 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -29,5 +29,8 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): time.sleep(3) plt.close('all') + # Simply test that this does not fail + results = model.save_results(dataset) + # If this fails, the reconstruction got worse assert model.loss_history[-1] < 6.5