code: def get_reference_grid(
      image_size: tuple[int] | list[int],
  ) -> torch.Tensor:
      mesh_points = [torch.arange(0, dim) for dim in image_size]
      grid = torch.stack(
          meshgrid_ij(*mesh_points), dim=0
      )  # (spatial_dims, ...)
      return grid.to(dtype=torch.float)
  
imports: from __future__ import annotations
  import torch
  from torch import nn
  from torch.nn import functional
  from monai.networks.blocks.regunet_block import RegistrationDownSampleBlock, RegistrationExtractionBlock, RegistrationResidualConvBlock, get_conv_block, get_deconv_block
  from monai.networks.utils import meshgrid_ij
label_desc: Use image node_id 2 for calling the AffineHead obj's underlying member method code with example usage: obj.get_reference_grid(image_size) and returns: grid