Organize the transforms so that they are easy to call in a machine learning pipeline. Right now it is a little cumbersome.
Idea is to make the transform a nn.Module so lightning fabric automatically set it up the devices can deal with it.
This approach also follows PyTorch conventions for transforms.
Make the directions a parameterless parameter so torch can cast it to the device.
Organize the transforms so that they are easy to call in a machine learning pipeline. Right now it is a little cumbersome.
Idea is to make the transform a nn.Module so lightning fabric automatically set it up the devices can deal with it.
This approach also follows PyTorch conventions for transforms.
Make the directions a parameterless parameter so torch can cast it to the device.