Dongkai Wang

Hello, I have created new commits to add the forward test, moving the losses and use pre-commit to check code styles.

This may be caused by PIL version. You can refer to this blog. https://discuss.pytorch.org/t/solved-typeerror-rotate-got-an-unexpected-keyword-argument-fillcolor/68115

Sorry I did not encounter this problem (the code has been tested on many different computers), maybe you should debug to see where the NaN appear.