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
To reproduce and see my proposed fix take a look at the following colab notebook:
@patch def lighting(x: TensorImage, func): return TensorImage(torch.sigmoid(func(logit(x))))
@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