AssertionError: Key "nms" missing from results in prediction script​ and GPU Out-of-Memory (OOM) Errors During Training

[13512263278]

Describe the bug

1. I'm using the code from the stadist-conic-2022 branch for training. However, after completing the training, when I run the prediction code, I encounter an error: assert "nms" in res. How can I resolve this issue?
2. My model trains successfully for over 300 epochs, but then suddenly crashes with CUDA out-of-memory errors. The failure occurs at different epochs each time

To reproduce
if refine_shapes is not False:
assert "nms" in res
u = refine(u, res, **refine_shapes)

Environment (please complete the following information):

• StarDist version 0.9.1
• CSBDeep version 0.8.1
• TensorFlow version 2.18.1
• OS: Ubuntu 22.04.2 LTS
• GPU memory (if applicable): 24GB

[13512263278]

@maweigert @uschmidt83 @ximion Sorry to bother you. Would you mind taking a look when you have a moment? Any thoughts or hints would be a great help.