code: def get_step_output_rows(
      self,
      state: State,
      unit: PredictUnit[TPredictData],
      # pyre-fixme[2]: Parameter annotation cannot be `Any`.
      step_output: Any,
  ) -> Union[List[str], List[List[str]]]:
      return [["1"], ["2"]]
  
file_path: tnt/tests/framework/callbacks/test_csv_writer.py
imports: import tempfile
  import unittest
  from typing import Any, List, Union
  from torchtnt.framework._test_utils import DummyPredictUnit, generate_random_dataloader
  from torchtnt.framework.callbacks.base_csv_writer import BaseCSVWriter
  from torchtnt.framework.predict import predict
  from torchtnt.framework.state import State
  from torchtnt.framework.unit import PredictUnit, TPredictData
label_desc: Use image node_id 1 for calling the CustomCSVWriter obj's underlying member method code with example usage: obj.get_step_output_rows(state, unit, step_output) and returns: list