Add torch no grad condition (#4391)
* Add torch no grad condition
Add torch no grad condition for param ds_tensor otherwise torch tracer will complain
* Update indentation
* Change Indentation to satisfy Yapf test
Change Indentation to satisfy Yapf test as the number of column lines was exceeding limit
* Change variable name to satisfy Yapf condition
Change variable name to satisfy Yapf condition otherwise the lenght was over 118 characters.
* remove trailing whitespace