From 7ca5bc59311cec237030104d3656c79137a068a1 Mon Sep 17 00:00:00 2001
From: Liu Liu <i@liuliu.me>
Date: Tue, 26 Nov 2024 17:57:43 -0500
Subject: [PATCH] Fix one more ModelBuilder / Model mismatch.

---
 nnc/ModelBuilder.swift | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/nnc/ModelBuilder.swift b/nnc/ModelBuilder.swift
index 769b52886f7..600b7a9b72b 100644
--- a/nnc/ModelBuilder.swift
+++ b/nnc/ModelBuilder.swift
@@ -121,17 +121,16 @@ public class AnyModelBuilder {
         { (handle, name, options, params, tensorOut) -> Int32 in
           let readerHelper = Unmanaged<DynamicGraph.Store.ModelReaderHelper>.fromOpaque(handle!)
             .takeUnretainedValue()
-          if tensorOut!.pointee == nil {
-            tensorOut!.pointee = ccv_nnc_tensor_new(nil, params, 0)
-          }
-          let cTensorOut = tensorOut!.pointee
-          let params = cTensorOut!.pointee.info
+          let params = tensorOut!.pointee?.pointee.info ?? params
           let result = readerHelper.reader(
             name.map { String(cString: $0) } ?? "", DataType.from(cTensorParams: params),
             TensorFormat.from(cTensorParams: params), TensorShape(dims: params.dim))
           switch result {
           case .final(let tensor):
             precondition(tensor.kind == .CPU)
+            if tensorOut!.pointee == nil {
+              tensorOut!.pointee = ccv_nnc_tensor_new(nil, params, 0)
+            }
             var input: UnsafeMutablePointer<ccv_nnc_tensor_t>? = tensor.cTensor
             ccv_nnc_cmd_exec(
               ccv_nnc_cmd(