特征图提取patch

最近复现代码时遇到了一些问题，原文中是先通过卷积网络提取特征以后，再把特征划分成一个个的patch，就比如以下代码，我想把两次卷积以后的
feature map 均匀划分成16个大小为6411*20的patch，这我该怎么办
class Net(torch.nn.Module):
def init(self):
super(Net, self).init()
self.conv1 = torch.nn.Sequential(
torch.nn.Conv2d(1, 10, kernel_size=5),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size=2),
)
self.conv2 = torch.nn.Sequential(
torch.nn.Conv2d(10, 20, kernel_size=5),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size=2),
)
self.fc = torch.nn.Sequential(
torch.nn.Linear(320, 50),
torch.nn.Linear(50, 10),
)

def forward(self, x):
    print("0",x.shape) #64,1,28,28
    batch_size = x.size(0)
    x = self.conv1(x)  # 64,10,12,12
    print ("1", x.shape)
    x = self.conv2(x)  # 64,20,4,4
    print("2",x.shape)
   


    # # for i in range[64]:
    # img = pathes        # 接上面操作，看shape变化
    # img = reduce (img, 'b n p2 -> n p2', 'mean')
    # print (img.shape)  # torch.Size([4, 196608])
    #
    # img = rearrange (img, 'n (p1 p2 c) -> n p1 p2 c', p1=2, p2=2, c=20)
    # print (img.shape)  # torch.Size([4, 256, 256, 3])
    #
    # # 画图，因为我是把图像ndarray转成tensor操作的，
    # # 所以要显示，就把它转换回numpy
    # fig, ax = plt.subplots (2, 2, figsize=(5, 5))
    # ax[0][0].imshow (img[0].data.numpy (), cmap='gray')
    # ax[0][1].imshow (img[1].data.numpy (), cmap='gray')
    # ax[1][0].imshow (img[2].data.numpy (), cmap='gray')
    # ax[1][1].imshow (img[3].data.numpy (), cmap='gray')
    # plt.show ()

    # x = x.view(batch_size, -1)  # flatten 变成全连接网络需要的输入 (batch, 20,4,4) ==> (batch,320), -1 此处自动算出的是320
    #
    # print(epoch,x.shape)
    # x = self.fc(x)
    return x