Many Thanks for your help! When i run the training model the file is saving the tensor which consist of 16 values. However the values seem to get over written at each step and the final CSV final only has 3 values. Any idea why I am not getting all the values? Thanks

I would guess that np.save() overrides the file every time and that the last batch from the dataloader is of size 3 (because dataset_len % batch_size = 3) and so you get this.
I am not sure if there is an option to append in np.save but a solution is to accumulate all the values in a Tensor and write it at the end.
Or write to the file yourself and make sure to append to it.