Which PyTorch version are you using?
If you are using 1.5, could you please update, as assert statements in 1.5.0 were not working properly.
Also, could you post the complete stack trace you are seeing?