Data file checkpoints

fix: Fixed training procedure

Detaching the output tensor ended up in a training failure.