Solved by albanD in post #2 
              
              
                Setting set_detect_anomaly doesn’t give any other output 

Make sure to use latest pytorch as we recently fix warnings not showing up in colab. 
Or run your code in command line to have the corresponding forward code. 
The code does quite a lot of inplace and viewing ops. 

Plus, it is working with…