[FeatureRequest] Generalize affine and grid transforms to support keys with different spatial size

[mibaumgartner]

Description
The grid is only created for the first element in the batch. This behaviour should be generalised to support keys with different spatial size without introducing computation overhead (e.g. computing and augmenting a grid multiple times even though the keys have the same spatial size)

it should be save to ignore the number of channels and only focus on spatial size of the grid because pytorch does not use that anyway (even though affine_grid wants the number of channels) https://github.com/pytorch/pytorch/blob/74b65c32be68b15dc7c9e8bb62459efbfbde33d8/aten/src/ATen/native/AffineGridGenerator.cpp#L34-L62