Code Cleanup (Ready to merge)

[collin-k]

Description
This PR addresses code readability and maintainability.

Changes

1. The one_trial function and the signature for the train function have been revised so that variables are passed as flat parameters and not as a part of a tuple.
2. The train_epoch function and the signature for the train_setup function have been revised so that variables are passed as flat parameters and not as a part of a tuple.
3. The train function and the signature for the test function have been revised so that variables are passed as flat parameters and not as a part of a tuple.

[annawoodard]

Great thank you!