Fixes CPU and hanging GPU crash (#2118)
Fixes CPU and hanging GPU crash (#2118)

* training batch clean up

* training batch clean up

* training batch clean up
williamFalcon/pytorch-lightning,williamFalcon/pytorch-lightning