code: def __init__(self, *, dim, patch_size, channels=3):
      super().__init__()
      patch_dim = patch_size * patch_size * 5 * channels
  
      self.to_patch_tokens = nn.Sequential(
          Rearrange(
              "b c (h p1) (w p2) -> b (h w) (p1 p2 c)",
              p1=patch_size,
              p2=patch_size,
          ),
          nn.LayerNorm(patch_dim),
          nn.Linear(patch_dim, dim),
      )
  
imports: from math import sqrt
  import torch
  import torch.nn.functional
  from torch import nn
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 1 to create a new SPT object from inherited base classes: nn with example: obj = SPT()