title: Extract weights of a tflite model directly on an android studio variable
body: <p>I'm having some difficulties writing an extract_weights and an initialize function for a tf.Module model that i later convert to tflite.</p>
  <p>The idea is that, i want to use this model for on device training.
  The project architecture is as it follows:
  -first i create a transfer learning model that will later be used for training
  -then i upload this model in my android application where i train it using the tflite.Interpreter
  -the model will be trained federated using a flower server
  The problem that i have at the moment is that flower needs to colect from each device the weights as ByteBuffers after each training loop, but i don't seem to understand how i could save them in my android application.</p>
  <p>These are the methods that i wrote</p>
  <pre><code>@tf.function
  def extract_weights(self):
      &quot;&quot;&quot;
      Extracts the traininable weights of the head model as a list of numpy arrays.
  
      Paramaters:
  
      Returns:
          Map of extracted weights and biases.
      &quot;&quot;&quot;
      tmp_dict = {}
      tensor_names = [weight.name for weight in self.head_model.weights]
      tensors_to_save = [weight.read_value() for weight in self.head_model.weights]
      for index, layer in enumerate(tensors_to_save):
          tmp_dict[tensor_names[index]] = layer
  
      return tmp_dict
  
  
  @tf.function(input_signature=[SIGNATURE_DICT])
  def initialize_weights(self, weights):
      &quot;&quot;&quot;
      Initializes weights of the head model.
  
      Paramaters:
          weights : Tensors used for initialization.
      Returns:
          NONE
      &quot;&quot;&quot;
      tensor_names = [weight.name for weight in self.head_model.weights]
      for i, tensor in enumerate(self.head_model.weights):
          tensor.assign(weights[tensor_names[i]])
  </code></pre>
  <p>To notice that when i instantiate a TransferLearningModel(my model class that implements tf.Module) object and call these to functions i got no problems but when i try to convert them to tflite i get this error:</p>
  <p><code>   ValueError: Got a non-Tensor value</code>&lt;tf.Operation 'StatefulPartitionedCall' type=StatefulPartitionedCall&gt;<code>for key 'output_0' in the output of the function __inference_initialize_weights_8582 used to generate the SavedModel signature 'initialize'. Outputs for functions used as signatures must be a      ValueError: Got a non-Tensor value</code>&lt;tf.Operation 'StatefulPartitionedCall' type=StatefulPartitionedCall&gt;<code> for key 'output_0' in the output of the function __inference_initialize_weights_8582 used to generate the SavedModel signature 'initialize'. Outputs for functions used as signatures must be a single Tensor, a sequence of Tensors, or a dictionary from string to Tensor.</code></p>
  <p>I understand the error but i don t get why i have to return something when simply initializing the weights of my model.</p>