如何使用重叠将PyTorch张量分块到指定的存储桶大小？

import torch

def overlap_chunk(tensor, chunk_size, overlap):
    """
    将张量分块，相邻块之间有重叠区域。
    
    :param tensor: 输入张量
    :param chunk_size: 每个块的大小 (height, width)
    :param overlap: 相邻块之间的重叠区域大小
    :return: 分块后的张量列表
    """
    h, w = tensor.shape[-2:]
    ch, cw = chunk_size
    oh, ow = overlap
    
    chunks = []
    for y in range(0, h - ch + 1, ch - oh):
        for x in range(0, w - cw + 1, cw - ow):
            chunk = tensor[..., y:y+ch, x:x+cw]
            chunks.append(chunk)
    
    return chunks

# 示例使用
tensor = torch.randn(1, 3, 64, 64)  # 假设是一个64x64的图像张量
chunk_size = (32, 32)
overlap = 8

chunks = overlap_chunk(tensor, chunk_size, overlap)
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}: {chunk.shape}")