Ming Li
change inference_mode to no_grad to avoid errors
acbc9f8