fix

fifield · Sep 9, 2024 · f92b1d9 · f92b1d9
1 parent 5e7318a
commit f92b1d9
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/lib/CAPI/Translation.cpp b/lib/CAPI/Translation.cpp
@@ -197,9 +197,13 @@ DEFINE_C_API_PTR_METHODS(AieRtxControl, xilinx::AIE::AIERTXControl)
 
 AieRtxControl getAieRtxControl(size_t partitionStartCol,
                                size_t partitionNumCols) {
-  NPUTargetModel targetModel;
-  AIERTXControl *ctl =
-      new AIERTXControl(targetModel);
+  std::vector<AIEDevice> devices{AIEDevice::npu1_1col, AIEDevice::npu1_2col,
+                                 AIEDevice::npu1_3col, AIEDevice::npu1_4col,
+                                 AIEDevice::npu1};
+  const BaseNPUTargetModel &targetModel =
+      (const BaseNPUTargetModel &)xilinx::AIE::getTargetModel(
+          devices[partitionNumCols - 1]);
+  AIERTXControl *ctl = new AIERTXControl(targetModel);
   return wrap(ctl);
 }