[AIR] Allow users to pass Callable[[torch.Tensor], torch.Tensor] to TorchVisionTransform

[bveeramani]

Signed-off-by: Balaji Veeramani balaji@anyscale.com

Why are these changes needed?

Transforms like RandomHorizontalFlip expect Torch tensors as input, but if you're applying the transform per-epoch, then you can't use ToTensor. To fix the problem, this PR updates TorchVisionPreprocessor to convert ndarray inputs to Torch tensors.

You can't use ToTensor to convert the ndarrays to Torch tensors because then you'd be applying ToTensor twice, and your images would get scaled incorrectly.

transform = ToTensor()
preprocessor = TorchVisionPreprocessor(["image"], transform=transform)

# You can't use `ToTensor` twice, because then images would be scaled incorrectly
per_epoch_transforms = Compose([ToTensor(), RandomHorizontalFlip(0.5)])
per_epoch_preprocessor = TorchVisionPreprocessor(["image"], transforms=per_epoch_transform) 

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[amogkam]

lgtm!

[bveeramani]

Test failures are unrelated.

[amogkam]

Can you please add a pr description?