class_docstr: Support TIMM image backbones.
  Refer to https://github.com/rwightman/pytorch-image-models
docstr: Parameters
  ----------
  images : torch.FloatTensor
      A tensor in [N, C, H, W] layout to represent the images.
  image_valid_num : torch.Tensor
      A tensor that describes valid number of input images.
  image_column_names : list of str, optional
      A list of strings that indicates names of the image columns.
  image_column_indices : list of torch.Tensor, optional
      A list of tensors that indicates start and stop indices of the image columns.
  
  Returns
  -------
      A dictionary with logits and features.
docstr_tok: ["Parameters","--","--","--","--","--","images",":","torch.FloatTensor","A","tensor","in","[","N",",","C",",","H",",","W","]","layout","to","represent","the","images",".","image_valid_num",":","torch.Tensor","A","tensor","that","describes","valid","number","of","input","images",".","image_column_names",":","list","of","str",",","optional","A","list","of","strings","that","indicates","names","of","the","image","columns",".","image_column_indices",":","list","of","torch.Tensor",",","optional","A","list","of","tensors","that","indicates","start","and","stop","indices","of","the","image","columns",".","Returns","--","--","--","-","A","dictionary","with","logits","and","features","."]
code: def forward(
      self,
      images: torch.FloatTensor,
      image_valid_num: torch.Tensor,
      image_column_names: Optional[List[str]] = None,
      image_column_indices: Optional[List[torch.Tensor]] = None,
  ):
      """
      Parameters
      ----------
      images : torch.FloatTensor
          A tensor in [N, C, H, W] layout to represent the images.
      image_valid_num : torch.Tensor
          A tensor that describes valid number of input images.
      image_column_names : list of str, optional
          A list of strings that indicates names of the image columns.
      image_column_indices : list of torch.Tensor, optional
          A list of tensors that indicates start and stop indices of the image columns.
  
      Returns
      -------
          A dictionary with logits and features.
      """
      if self.mix_choice == "all_images":  # mix inputs
          mixed_images = (
              images.sum(dim=1)
              / torch.clamp(image_valid_num, min=1e-6)[
                  :, None, None, None
              ]
          )  # mixed shape: (b, 3, h, w)
          features = self.model(mixed_images)
          if self.num_classes > 0:
              logits = self.head(features)
          else:
              logits = features
  
          column_features = {}
          column_feature_masks = {}
  
      elif self.mix_choice == "all_logits":  # mix outputs
          b, n, c, h, w = images.shape
          features = self.model(
              images.reshape((b * n, c, h, w))
          )  # (b*n, num_features)
          if self.num_classes > 0:
              logits = self.head(features)
          steps = torch.arange(0, n).type_as(image_valid_num)
          image_masks = (
              steps.reshape((1, -1)) < image_valid_num.reshape((-1, 1))
          ).type_as(
              features
          )  # (b, n)
          features = (
              features.reshape((b, n, -1)) * image_masks[:, :, None]
          )  # (b, n, num_features)
  
          batch = {
              self.image_key: images,
              self.image_valid_num_key: image_valid_num,
          }
          if image_column_names:
              assert len(image_column_names) == len(
                  image_column_indices
              ), "invalid image column inputs"
              for idx, name in enumerate(image_column_names):
                  batch[name] = image_column_indices[idx]
  
          # collect features by image column names
          column_features, column_feature_masks = get_column_features(
              batch=batch,
              column_name_prefix=self.image_column_prefix,
              features=features,
              valid_lengths=image_valid_num,
          )
  
          features = (
              features.sum(dim=1)
              / torch.clamp(image_valid_num, min=1e-6)[:, None]
          )  # (b, num_features)
          if self.num_classes > 0:
              logits = (
                  logits.reshape((b, n, -1)) * image_masks[:, :, None]
              )  # (b, n, num_classes)
              logits = (
                  logits.sum(dim=1)
                  / torch.clamp(image_valid_num, min=1e-6)[:, None]
              )  # (b, num_classes)
          else:
              logits = features
  
      else:
          raise ValueError(f"unknown mix_choice: {self.mix_choice}")
  
      if column_features == {} or column_feature_masks == {}:
          return features, logits
      else:
          return features, logits, column_features, column_feature_masks
  
imports: import json
  import logging
  import os
  from typing import Dict, List, Optional
  import torch
  from timm import create_model
  from timm.layers.linear import Linear
  from torch import nn
  from ..constants import AUTOMM, COLUMN, COLUMN_FEATURES, FEATURES, IMAGE, IMAGE_VALID_NUM, LABEL, LOGITS, MASKS
  from .utils import assign_layer_ids, get_column_features, get_model_head
label_desc: Use image node_id 9 for calling the TimmAutoModelForImagePrediction obj's underlying member method code with example usage: obj.forward(images, image_valid_num, image_column_names, image_column_indices) and returns: features,  logits, features,  logits,  column_features,  column_feature_masks