I would look at the vision transformer architecture (so e.g. DeiT-Tiny to name one).
They have an embedding layer, too, that splits the images into 16x16 blocks and applies a single convolution to many channels to then add a learned positional embedding.
You could use the output of the convolutional part as the blocks (so skipping that conv in the vision embedding) and just the learned embeddings.
The vision transformers also add a (learned) global point (the class token) to the patches. This is where they then hook the prediction head to at the end of the transformer. This class token originates from BERT and, combined with its attention, essentially replaces the pooling to go from patches to global information.

  

      github.com
  

  
    MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L5-L28



      
      
          class VisionPatchEmbedding(torch.nn.Module):
              def __init__(self, input_size, patch_size, n_channels):
                  super().__init__()
                  self.n_channels = n_channels
                  num_patches = (input_size ** 2) // (
                      patch_size ** 2
                  )  # only when this is divisible...
                  self.patch_to_vec = torch.nn.Conv2d(
                      3,
                      n_channels,
                      kernel_size=(patch_size, patch_size),
                      stride=(patch_size, patch_size),
                  )  # equivalent to reshape + linear
                  self.class_token = torch.nn.Parameter(
                      torch.randn(1, 1, n_channels)
                  )  # "global information"
                  self.pos_embedding = torch.nn.Parameter(
                      torch.randn(1, num_patches + 1, n_channels)
                  )
          
      

    
  

  This file has been truncated. show original

  

  
    
    
  

  


The file has a DeiTTiny class a bit further down.
Best regards
Thomas