Remove side effects of disabling gradient computaiton
Disabling gradient computation here will affect all next operations using torch, and will remove the gradient computation. This made a few `Trainer` test fail in slow tests.