Function affine_grid_2d

pub fn affine_grid_2d<B>(
    transform: Tensor<B, 3>,
    dims: [usize; 4],
) -> Tensor<B, 4>
where
    B: Backend,

Generate a tensor with homogeonous coordinates of each element’s transformed location

See:

• torch.nn.functional.affine_grid

• transform - Transformation with shape (batch_size, 2, 3)
• dims - dimensions as (batch_size, channels, height, width)

Returns

Tensor with shape (batch_size, height, width, 2), where dim 2 is (x, y) All coordinates are broadcast on the batch dim