docstr: Get the loss data from FLServer and construct the identical Pytorch Tensor
docstr_tok: ["Get","the","loss","data","from","FLServer","and","construct","the","identical","Pytorch","Tensor"]
code: def train_step(self, x, y):
      """
      Get the loss data from FLServer and construct the identical Pytorch Tensor
      """
      with tf.GradientTape() as tape:
          y_pred_local = self.model(x)
      y_true_global = y
      data_map = {"input": y_pred_local.numpy()}
      if y_true_global is not None:
          data_map["target"] = y_true_global.numpy()
      logging.debug(
          f"[{self.fl_client.client_uuid}] client sending train data to server"
      )
      response = self.fl_client.train(data_map)
      logging.debug(
          f"[{self.fl_client.client_uuid}] client got response from server"
      )
      response_map = tensor_map_to_ndarray_map(response.data.tensorMap)
      server_grad = response_map["grad"]
      gradients = tape.gradient(
          y_pred_local,
          self.model.trainable_variables,
          tf.convert_to_tensor(server_grad),
      )
      self.optimizer.apply_gradients(
          zip(gradients, self.model.trainable_variables)
      )
      return response_map["loss"]
  
imports: import logging
  import os
  import shutil
  import tempfile
  from numpy import ndarray
  from bigdl.dllib.utils.log4Error import invalidInputError
  from bigdl.ppml.fl.nn.fl_client import FLClient
  from bigdl.ppml.fl.nn.nn_client import NNClient
  from bigdl.ppml.fl.nn.utils import file_chunk_generate, print_file_size_in_dir, tensor_map_to_ndarray_map
  import tensorflow
  from tensorflow.keras.layers import Dense, Flatten, Conv2D
  from tensorflow.keras import Model
  from bigdl.ppml.fl.nn.generated.nn_service_pb2 import ByteChunk
  from nn_service_pb2 import LoadModelRequest, SaveModelRequest
label_desc: Use image node_id 6 for calling the TensorflowEstimator obj's underlying member method code with example usage: obj.train_step(x, y) and returns: response_map