【Pytorch Function】torch.meshgrid用法

torch.meshgrid

Take N tensors, each of which can be either scalar or 1-dimensional vector, and create N N-dimensional grids, where the i-th grid is defined by expanding the i-th input over dimensions defined by other inputs.

x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5, 6])
grid_x, grid_y = torch.meshgrid(x, y)

grid_x

tensor([[1, 1, 1],
[2, 2, 2],
[3, 3, 3]])

grid_y

tensor([[4, 5, 6],
[4, 5, 6],
[4, 5, 6]])

参考