title: Pre-training weights with one loss function and then change loss function to a custom one
body: <p>I'm using the layers API and i want to pre-train my weights in order to create better initialisation values for a loss function that has some requirements about the output values from the network.</p>
  
  <p>I'm pre-training atm on some generated data that has random values inside the ranges i want the output to me. It is trained by just using a MSE loss.</p>
  
  <p>However after I pre-trained the network i would want to change the loss function and of cause also the labels however the labels are not a problem. When I change te loss function and try to run the network again I get an error:</p>
  
  <pre><code>InvalidArgumentError: TensorArray has inconsistent shapes.  
  Index 0 has shape: [1] but index 1 has shape: []
       [[Node: map/TensorArrayStack/TensorArrayGatherV3 =
  TensorArrayGatherV3[_class=["loc:@map/TensorArray_1"], 
  dtype=DT_FLOAT, element_shape=&lt;unknown&gt;, _device="/job:localhost/replica:0/task:0/cpu:0"]
  (map/TensorArray_1, map/TensorArrayStack/range, map/while/Exit_1)]]
  </code></pre>
  
  <p>I don't know if it is impossible to switch loss functions due to some tensors in Adam maybe needs to be another shape or something.</p>
  
  <p>If anyone knows please enlighten me.</p>
  
  <p><strong>Custom Loss:</strong></p>
  
  <pre><code>def total_loss(y_true, y_pred):
  
      def single_loss(total_index):
  
          def iou_fn(box, y_true_x, y_true_w):
              box_x, box_w = box[0], box[1]
              left = tf.maximum(y_true_x-y_true_w/2., box_x-box_w/2.)
              right = tf.minimum(y_true_x+y_true_w/2., box_x+box_w/2.)
              overlap = right - left
  
              intersection = tf.maximum(overlap*1., 0.)
  
              union = y_true_w*1. + box_w*1. - intersection
  
              return (intersection*1.) / (union*1.)
  
          cell_size = 16
          box_count = 5
          cell_count = 10
  
          y_true_single = y_true[total_index]
          y_pred_single = y_pred[total_index]
  
          coord = tf.constant(5., name='coord')
          noobj = tf.constant(0.5, name='noobj')
  
          x_error = 0.
          w_error = 0.
          c_error = 0.
          c_noobj_error = 0.
          class_error = 0.
  
          y_true_x, y_true_w, y_true_class = tf.split(y_true_single[1:4], 3, axis=0)
  
          for cell_index in range(cell_count):
              cell = y_pred_single[cell_index*cell_size:(cell_index+1)*cell_size]
              boxes = cell[0:-1]
              _class = cell[-1]
              boxes_split = tf.split(boxes, box_count, axis=0)
              ious = []
              for box_index in range(box_count):
                  ious.append(iou_fn(boxes_split[box_index], y_true_x, y_true_w))
                  index = tf.argmax(ious, output_type=tf.int32)
              x_error = tf.cond((tf.gather(ious, tf.argmax(ious, output_type=tf.int32))[0]&gt;0.)[0], 
                                lambda: tf.add(x_error, (tf.square(tf.gather(boxes_split, index)[0][0]-y_true_x[0]))), 
                                lambda: x_error)
  
              w_error = tf.cond((tf.gather(ious, tf.argmax(ious, output_type=tf.int32))[0]&gt;0.)[0], 
                                lambda: tf.add(w_error, (tf.square(tf.gather(boxes_split, index)[0][1]-y_true_w[0]))), 
                                lambda: w_error)
  
              c_error = tf.cond((tf.gather(ious, tf.argmax(ious, output_type=tf.int32))[0]&gt;0.)[0], 
                                lambda: tf.add(c_error, (tf.square(tf.gather(ious, index)[0][0]))), 
                                lambda: c_error)
  
              class_error = tf.cond((tf.gather(ious, index)[0]&gt;0.)[0], 
                                lambda: tf.add(class_error, tf.square(_class-y_true_class)), 
                                lambda: class_error)
  
              for box_index in range(box_count):
                  c_noobj_error = tf.cond((tf.gather(ious, index)[0]&gt;0.)[0], 
                                lambda: tf.cond(tf.equal(box_index, index[0]), 
                                                lambda: c_noobj_error, 
                                                lambda: tf.add(c_noobj_error, tf.square(tf.gather(ious, index)[0]))), 
                                lambda: c_noobj_error)
  
          loss = x_error*coord + w_error*coord + c_error + c_noobj_error*noobj + class_error
  
          return loss
  
      full_return = tf.map_fn(single_loss, tf.range(tf.shape(y_pred)[0]), dtype=tf.float32)
      return full_return
  </code></pre>