class_docstr: Manages a PyTorch model for training.
class_docstr_tok: ["Manages","a","PyTorch","model","for","training","."]
code: def train_epochs(
      self,
      data_creator,
      epochs=1,
      batch_size=32,
      profile=False,
      wrap_dataloader=None,
      callbacks=None,
      validation_data_creator=None,
  ):
      if self.state_dict:
          self.load_state_dict(self.state_dict.value)
      stats_list = super().train_epochs(
          data_creator=data_creator,
          epochs=epochs,
          batch_size=batch_size,
          profile=profile,
          wrap_dataloader=wrap_dataloader,
          callbacks=callbacks,
          validation_data_creator=validation_data_creator,
      )
  
      if self.log_to_driver:
          LogMonitor.stop_log_monitor(
              self.log_path, self.logger_thread, self.thread_stop
          )
  
      if self.model_dir is not None:
          if self.rank == 0:
              state_dict = self.get_state_dict()
              save_pkl(
                  state_dict, os.path.join(self.model_dir, "state.pkl")
              )
          return [stats_list]
      else:
          if self.rank == 0:
              state_dict = self.get_state_dict()
              return [state_dict, stats_list]
          else:
              return [stats_list]
  
file_path: BigDL/python/orca/src/bigdl/orca/learn/pytorch/pytorch_pyspark_worker.py
imports: from bigdl.orca.learn.pytorch.torch_runner import TorchRunner
  import torch.distributed
  import logging
  import os
  import tempfile
  import copy
  from pyspark import BarrierTaskContext, TaskContext
  from bigdl.orca.learn.utils import save_pkl, duplicate_stdout_stderr_to_file, get_rank, get_partition_id
  from bigdl.orca.learn.log_monitor import LogMonitor
label_desc: Use image node_id 4 for calling the PytorchPysparkWorker obj's underlying member method code with example usage: obj.train_epochs(data_creator, epochs, batch_size, profile, wrap_dataloader, callbacks, validation_data_creator) and returns: list, list, list