-use range instead of a tensor with the same values
-small refactoring to make the code easier to read.