如何在PyTorch中将坐标值的3D网格分割成块？

import torch

# 定义网格块的大小
block_size = 2

# 创建网格坐标
x = torch.arange(0, 6)
y = torch.arange(0, 6)
z = torch.arange(0, 6)
grid = torch.meshgrid(x, y, z)

# 分割网格坐标
blocks = []
for i in range(0, len(x), block_size):
    for j in range(0, len(y), block_size):
        for k in range(0, len(z), block_size):
            block = grid[0][i:i+block_size, j:j+block_size, k:k+block_size]
            blocks.append(block)

# 转换块的表示
for block in blocks:
    # 根据块的网格坐标从原始3D网格中选择对应的块
    selected_block = torch.index_select(original_grid, 0, block.flatten())
    # 进行块的处理操作
    # ...

# 输出每个块的网格坐标
for block in blocks:
    print(block)