@pmeier Looks good. Let's do the same for AA.

Not sure why, but it seems AA seems to work as is:

import pickle

import torch
from torchvision.prototype import transforms


image = torch.randint(0, 256, (3, 512, 512), dtype=torch.uint8)

for transform in [
    transforms.AutoAugment(),
    transforms.AugMix(),
    transforms.AutoAugment(),
    transforms.RandAugment(),
    transforms.TrivialAugmentWide(),
]:
    serialized = pickle.dumps(transform)
    deserialized = pickle.loads(serialized)

    deserialized(image)

I suggested we don't make it more complicated until we actually hit a problem. Using the same technique that I used here, is not really possible for AA, since the return value is dependent on the input parameters.

Using the same technique that I used here, is not really possible for AA, since the return value is dependent on the input parameters.

Isn't the return type of the lambda always Optional[Tensor]?

Isn't the return type of the lambda always Optional[Tensor]?

Yeah, but the return value is different and depends on the inputs. functools.partial only works in this PR, because the return type is static and we know it in advance. For AA we would need to replace

with something like

def _auto_augment_shearx(num_bins: int, height: int, width: int) -> torch.Tensor:
    return torch.linspace(0.0, 0.3, num_bins)

If we would know num_bins at class definition, we could do something like

functools.partial(torch.linspace, 0.0, 0.3, STATIC_NUM_BINS)

This is indeed fixed for some of the AA transforms like

but RandAugment for example, takes this as input parameter:

Worse, although not used often, know anything about height and width at definition time. Thus, all definitions that use height or width would need to have a standalone function.