:bug: Fix loss tensor to float conversion

[gkumbhat]

Changes

• Change how we fetch the float value from loss tensor to avoid potentially weird memory issues

Ref blog talking about this: https://medium.com/@raghadalghonaim/memory-leakage-with-pytorch-23f15203faa4