"Cell Detection with Contour Proposal Networks" notebook crashes on forward pass

[codingS3b]

I'm having trouble running through your demo notebook "Cell Detection with Contour Proposal Networks.ipynb". The kernel crashes when model training is started on the first forward pass. I narrowed it down to the execution of the forward function of the self.body inside the BackboneAsUNet object. I suspected GPU memory issues and lowered the batch size to 1 in the config and the crop_size to (128, 128) but without success.

I'm running on an Nvidia A100 (40GB), so I would really be surprised about a memory problem. Python 3.11.9 is used with torch 2.3.0+cu118 (installed via pip).

Would be happy to get any hints on the origin of the problem.