BUG (?): lightning patch overwrites type

As requested from the GitHub issue template, I am going to post here first.

See the following line in the vision augmentation code:

The lightning transform is patched to all TensorImage types but converts the input to a TensorImage and breaks sub-classed TensorImage code.

To reproduce and see my proposed fix take a look at the following colab notebook:

TL;DR:

@patch
def lighting(x: TensorImage, func): return TensorImage(torch.sigmoid(func(logit(x))))

to

@patch
def lighting(x: TensorImage, func): return (torch.sigmoid(func(logit(x)))

But I may be missing why the conversion to TensorImage is necessary.
For me at least, it is unexpected behavior and breaks my subclass.

Would be more than happy to create a tiny PR and hear some comments :slight_smile: