Update (base update)

[ghstack-poisoned]

Author: JCNTH

backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp

@@ -101,9 +101,18 @@ void fused_ce(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   const ValueRef logits = args[arg_idx++];
   const ValueRef labels = args[arg_idx++];
   const ValueRef n_valid_ref = args[arg_idx++];
-  const ValueListPtr out_tuple = graph.get_value_list(args[arg_idx++]);
-  const ValueRef loss = out_tuple->at(0);
-  const ValueRef dlogits = out_tuple->at(1);
+  const ValueRef out_tuple_ref = args[arg_idx++];
+
+  // Release the ValueListPtr before allocating loss_partial: add_tensor (via
+  // TmpTensor) asserts no live get_*() pointer exists, because adding Values
+  // may reallocate values_ and invalidate it (check_no_active_value_ptrs).
+  ValueRef loss = kDummyValueRef;
+  ValueRef dlogits = kDummyValueRef;
+  {
+    const ValueListPtr out_tuple = graph.get_value_list(out_tuple_ref);
+    loss = out_tuple->at(0);
+    dlogits = out_tuple->at(1);
+  }
 
   VK_CHECK_COND(
       graph.is_buffer_storage(logits) && graph.is_buffer_storage(dlogits),