class_docstr: Wraps model forward pass + predictions into a single nn.Module.
  
  The forward call of this module returns a flattened dictionary in order to support Triton input/output.
  
  TODO(geoffrey): Implement torchscript-compatible feature_utils.LudwigFeatureDict to replace
  get_module_dict_key_from_name and get_name_from_module_dict_key usage.
class_docstr_tok: ["Wraps","model","forward","pass","+","predictions","into","a","single","nn.Module",".","The","forward","call","of","this","module","returns","a","flattened","dictionary","in","order","to","support","Triton","input\/output",".","TODO","(","geoffrey",")",":","Implement","torchscript-compatible","feature_utils.LudwigFeatureDict","to","replace","get_module_dict_key_from_name","and","get_name_from_module_dict_key","usage","."]
code: def forward(
      self, preproc_inputs: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      model_outputs = self.model(preproc_inputs)
      predictions_flattened: Dict[str, torch.Tensor] = {}
      for module_dict_key, predict in self.predict_modules.items():
          feature_name = get_name_from_module_dict_key(module_dict_key)
          feature_predictions = predict(model_outputs, feature_name)
          # Flatten out the predictions to support Triton input/output
          for predict_key, tensor_values in feature_predictions.items():
              predict_concat_key = (
                  output_feature_utils.get_feature_concat_name(
                      feature_name, predict_key
                  )
              )
              predictions_flattened[predict_concat_key] = tensor_values
      return predictions_flattened
  
imports: import logging
  import os
  from typing import Any, Dict, Optional, TYPE_CHECKING, Union
  import pandas
  import torch
  from torch import nn
  from ludwig.constants import NAME, POSTPROCESSOR, PREDICTOR, PREPROCESSOR, TYPE
  from ludwig.data.postprocessing import convert_dict_to_df
  from ludwig.data.preprocessing import load_metadata
  from ludwig.features.feature_registries import get_input_type_registry
  from ludwig.features.feature_utils import get_module_dict_key_from_name, get_name_from_module_dict_key
  from ludwig.globals import MODEL_HYPERPARAMETERS_FILE_NAME, TRAIN_SET_METADATA_FILE_NAME
  from ludwig.types import ModelConfigDict, TrainingSetMetadataDict
  from ludwig.utils import output_feature_utils
  from ludwig.utils.data_utils import load_json, save_json
  from ludwig.utils.inference_utils import get_filename_from_stage, to_inference_module_input_from_dataframe
  from ludwig.utils.misc_utils import get_from_registry
  from ludwig.utils.output_feature_utils import get_feature_name_from_concat_name, get_tensor_name_from_concat_name
  from ludwig.utils.torch_utils import DEVICE
  from ludwig.utils.types import TorchDevice, TorchscriptPreprocessingInput
label_desc: Use image node_id 2 for calling the _InferencePredictor obj's underlying member method code with example usage: obj.forward(preproc_inputs) and returns: predictions_flattened