Is your feature request related to a problem? Please describe.
Feature from internal user:
Hello. I'm working with generative models for 3D scans (size ~200^3). Due to the cost of performing the RandAffine transform on cpu, the gpus are often waiting for the data loader.
create_grid takes up a large fraction of the time. By caching the cartesian grid between iterations (and creating copies when applying the transform) we can reduce the cost of RandAffine on cpu by 4 (0.55 to 0.13 sec) and on gpu by 100 (0.45 to 0.0045 sec). This change should allow the client to scale up the inputs when needed.
Caching is not particularly elegant tho. Would a pull request along these lines be acceptable?

The change would be something like adding this method:

def create_cartesian_grid(
        self,
        sp_size: Union[Sequence[float], float]
    ) -> Union[np.ndarray, torch.Tensor]:
        if self.use_cached_grid:
            if self.cached_grid is None or not sp_size == self.cached_sp_size:
                self.cached_sp_size = sp_size
                self.cached_grid = create_grid(spatial_size=sp_size)
                if not torch.is_tensor(self.cached_grid):
                    self.cached_grid = torch.tensor(self.cached_grid)
                if self.device:
                    self.cached_grid = self.cached_grid.to(self.device)
            return self.cached_grid
        else:
            return create_grid(spatial_size=sp_size)

In this case, self.use_cached_grid would be given by argument in the constructor and if False the behavior would be same as now.