Function affine_grid_2d 

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

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