@smth I managed to pinpoint the actual cause of the issue. It was a silly mistake in my code. In the loader thread, I was loading the ground truth as cuda tensor.   This was causing the deadlock when the main thread was holding the lock on GPUs while training. Now I fixed it.