fixes in rotate and reflect functions

[sridhs21]

This PR fixes incorrect tensor dimension handling in the reflect() function, improves mask interpolation in the rotate() function to preserve binary nature of segmentation masks, and removes redundant train_loss and val_loss declarations that reset loaded loss history.

Problems:

• reflect() function was mixing 2D and 3D tensor operations which caused incorrect flips for all tensor (containing 4 channels of physical fields).
• Using bilinear interpolation for binary masks in the rotate() function introduced intermediate values, which could corrupt the binary ground truth (i.e. instead of 0, 1, we can get 0.25, 0.75).
• Extra declaration of train_loss and val_loss resetting loss history.

Changes made:

• Changed BILINEAR to NEAREST for mask in rotate function to preserve binary values (0s and 1s only).
• Corrected reflection operations to use consistent 3D tensor dimensions by changing dims=(axis,) to dims=(axis+1,) for all tensors, fixing the incorrect dimension mapping for 3D tensors.
• Removed extra train_loss and val_loss declarations after checkpoint loading that were resetting the loaded loss history.

[cwsmith]

Looks good. Thank you.