class_docstr: A modified conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a
  sample shaped output.
  
  This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
  for all models (such as downloading or saving).
docstr: Sets the attention processor to use [feed forward
  chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
  
  Parameters:
      chunk_size (`int`, *optional*):
          The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
          over each tensor of dim=`dim`.
      dim (`int`, *optional*, defaults to `0`):
          The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
          or dim=1 (sequence length).
docstr_tok: ["Sets","the","attention","processor","to","use","[","feed","forward","chunking","]","(","https",":","\/\/huggingface.co\/blog\/reformer","#","2-chunked-feed-forward-layers",")",".","Parameters",":","chunk_size","(","`","int","`",",","*","optional","*",")",":","The","chunk","size","of","the","feed-forward","layers",".","If","not","specified",",","will","run","feed-forward","layer","individually","over","each","tensor","of","dim=","`","dim","`",".","dim","(","`","int","`",",","*","optional","*",",","defaults","to","`","0","`",")",":","The","dimension","over","which","the","feed-forward","computation","should","be","chunked",".","Choose","between","dim=0","(","batch",")","or","dim=1","(","sequence","length",")","."]
code: def enable_forward_chunking(
      self, chunk_size: Optional[int] = None, dim: int = 0
  ) -> None:
      """
      Sets the attention processor to use [feed forward
      chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
  
      Parameters:
          chunk_size (`int`, *optional*):
              The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
              over each tensor of dim=`dim`.
          dim (`int`, *optional*, defaults to `0`):
              The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
              or dim=1 (sequence length).
      """
      if dim not in [0, 1]:
          raise ValueError(
              f"Make sure to set `dim` to either 0 or 1, not {dim}"
          )
  
      # By default chunk size is 1
      chunk_size = chunk_size or 1
  
      def fn_recursive_feed_forward(
          module: torch.nn.Module, chunk_size: int, dim: int
      ):
          if hasattr(module, "set_chunk_feed_forward"):
              module.set_chunk_feed_forward(
                  chunk_size=chunk_size, dim=dim
              )
  
          for child in module.children():
              fn_recursive_feed_forward(child, chunk_size, dim)
  
      for module in self.children():
          fn_recursive_feed_forward(module, chunk_size, dim)
  
imports: from typing import Any, Dict, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.utils.checkpoint
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..loaders import UNet2DConditionLoadersMixin
  from ..utils import logging
  from .attention_processor import ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor
  from .embeddings import TimestepEmbedding, Timesteps
  from .modeling_utils import ModelMixin
  from .transformer_temporal import TransformerTemporalModel
  from .unet_2d_blocks import UNetMidBlock2DCrossAttn
  from .unet_2d_condition import UNet2DConditionModel
  from .unet_3d_blocks import CrossAttnDownBlockMotion, CrossAttnUpBlockMotion, DownBlockMotion, UNetMidBlockCrossAttnMotion, UpBlockMotion, get_down_block, get_up_block
  from .unet_3d_condition import UNet3DConditionOutput
label_desc: Use image node_id 8 for calling the UNetMotionModel obj's underlying member method code with example usage: obj.enable_forward_chunking(chunk_size, dim) without return types