code: def forward(
      self,
      image: torch.Tensor,
      timesteps: Optional[torch.Tensor] = None,
      return_probe_features: bool = False,
  ) -> torch.Tensor:
      n_batch = image.shape[0]
      h = self.blocks["input"](image, t=timesteps)
  
      for i in range(self.n_xf_blocks):
          h = self.blocks[f"block_{i}"](h)
  
      h = self.blocks["output"](
          h, return_probe_features=return_probe_features
      )
  
      assert list(h.shape) == [
          n_batch,
          self.n_embd if not return_probe_features else self.n_state,
      ]
  
      return h
  
imports: import math
  from collections import OrderedDict
  from typing import List, Optional, Tuple, cast
  import attr
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from .attention import AttentionInfo, DenseAttentionMask, DenseCausalAttentionMask, make_full_layout, to_attention_info
  from .utils import Affine, LayerNorm, zero_key_bias_grad
label_desc: Use image node_id 2 for calling the ImageEncoder obj's underlying member method code with example usage: obj.forward(image, timesteps, return_probe_features) and returns: h