[TODO]: Add unit tests for recently added *Transform classes

[hwpang]

Notes
See here for details. Copied the comment below:

LGTM. Perhaps we should use the *Transform classes in the integration tests to catch this type of error. That can be a separate PR.

good point! We should have:

• unit tests for our transforms: ensure the math/implementation is correct
• integration tests for our new models: ensure that training and evaluation output domains are different, i.e., our training outputs should have μ^≃0,σ^≃1 and our evaluation outputs should have μ^≃μ,σ^≃σ (roughly equal to our training statistics)