Pytorch two-layer loop to realize the segmentation of large pictures

The size of the large picture is resize(672,1120), Each piece 224x224

for i, (inputs, labels) in enumerate(self.train_data_loader):  

    crop_img=torch.zeros((5,3))
    for ii in range(5):
         for j in range(3):
             if ii==0 and j==0:
                  crop_img=inputs[:,:,224*j:224*(j+1),224*ii:224*(ii+1)]
             else:
                  crop_img=torch.cat((crop_img,inputs[:,:,224*j:224*(j+1),224*ii:224*(ii+1)]),0)

     grid_img = torchvision.utils.make_grid(crop_img.cpu(), nrow=16)
     plt.imshow(grid_img.permute(1, 2, 0))
     plt.show()