code: def forward(
      self, kjt: KeyedJaggedTensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      fp_ebc_out = self._fp_ebc(kjt)
      pred = torch.cat(
          [
              fp_ebc_out[key]
              for key in [
                  "feature_0",
                  "feature_1",
                  "feature_2",
                  "feature_3",
              ]
          ],
          dim=1,
      )
      loss = pred.mean()
      return loss, pred
  
file_path: torchrec/torchrec/distributed/tests/test_fp_embeddingbag_utils.py
imports: from typing import cast, Dict, List, Tuple
  import torch
  from torch import nn
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.feature_processor_ import FeatureProcessor, PositionWeightedModule, PositionWeightedModuleCollection
  from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the SparseArch obj's underlying member method code with example usage: obj.forward(kjt) and returns: loss,  pred