在 PyTorch 中，`chunk` 是一个用于将张量（tensor）按指定维度进行切片的函数。它可以将一个张量切分成多个块。
下面是一个使用 `chunk` 函数的示例：
```python
import torch
创建一个大小为 (6, 8) 的张量
tensor = torch.arange(48).reshape(6, 8)
print(tensor)
输出:
tensor([[ 0, 1, 2, 3, 4, 5, 6, 7],
[ 8, 9, 10, 11, 12, 13, 14, 15],
[16, 17, 18, 19, 20, 21, 22, 23],
[24, 25, 26, 27, 28, 29, 30, 31],
[32, 33, 34, 35, 36, 37, 38, 39],
[40, 41, 42, 43, 44, 45, 46, 47]])
使用 chunk 函数将张量在第 1 维度上切分成两个块
chunks = torch.chunk(tensor, 2, dim=0)
for chunk in chunks:
print(chunk)
输出:
tensor([[ 0, 1, 2, 3, 4, 5, 6, 7],
[ 8, 9, 10, 11, 12, 13, 14, 15],
[16, 17, 18, 19, 20, 21, 22, 23]])
tensor([[24, 25, 26, 27, 28, 29, 30, 31],
[32, 33, 34, 35, 36, 37, 38, 39],
[40, 41, 42, 43, 44, 45, 46, 47]])
```
在上面的例子中，我们首先创建了一个大小为 (6, 8) 的张量 `tensor`。然后，我们使用 `chunk` 函数将 `tensor` 在第 1 维度上切分成两个块。`chunk` 函数的第一个参数是要切分的张量，第二个参数是要切分的块数，第三个参数 `dim` 是指定切分的维度。
通过循环遍历 `chunks`，我们可以分别打印出切分后的两个块。可以看到，原始张量在第 1 维度上被均匀切分成两个大小相等的子张量。
`chunk` 函数在处理大型张量时非常有用，可以将其分割成更小的块，以便逐块处理或并行处理。