code: def forward(self, x):
      x = self.feature_extractor(x)
      x = torch.flatten(x, 1)
      logits = self.classifier(x)
      probs = torch.nn.functional.softmax(logits, dim=1)
      return logits, probs
  
file_path: DeepSpeed/tests/unit/profiling/flops_profiler/test_flops_profiler.py
imports: import torch
  import pytest
  import deepspeed
  from deepspeed.profiling.flops_profiler import get_model_profile
  from unit.simple_model import SimpleModel, random_dataloader
  from unit.common import DistributedTest
  from deepspeed.runtime.utils import required_torch_version
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 2 for calling the LeNet5 obj's underlying member method code with example usage: obj.forward(x) and returns: logits,  probs