Deadlock

[jeffzwe]

Running the package with my custom dataset results in a deadlock.

2024-06-18 11:52:18.824864: unpacking dataset...
2024-06-18 11:58:07.176028: unpacking done...
2024-06-18 11:58:07.179740: do_dummy_2d_data_aug: False
2024-06-18 11:58:07.419208: Unable to plot network architecture:
2024-06-18 11:58:07.419345: No module named 'hiddenlayer'
2024-06-18 11:58:07.825781:
2024-06-18 11:58:07.825900: Epoch 0
2024-06-18 11:58:07.826053: Current learning rate: 0.01
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
using pin_memory on device 0