Training hangs after a few epochs

Training would hang after a few epochs. If i do a keyboard interrupt and try to restart training it will hang right away without finishing a forward pass. Also Keyboard interrupt after it hangs doesn’t clear GPU memory so I have to restart the entire system to get the memory back. Restarting the system allows me to continue training.

Framework: torch==1.5.0
Windows: Version 2004 (OS Build 20161.1000)
WSL2 Ubuntu 20.04