Fix tensor transfer for dynamic shape specialization

iree-org · Oct 21, 2024 · 916ab3a · 916ab3a
1 parent 6d6f117
commit 916ab3a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/iree/turbine/ops/iree.py b/iree/turbine/ops/iree.py
@@ -83,8 +83,10 @@ class transfer_to_logical_device(CustomOp):
     def select(self, ksel: KernelSelection):
         ksel.attr_str(0)
         ta = ksel.arg_tensor(1)
-        ta.specialize_all_dims()
-        ksel.return_tensor(ta.t).specialize_all_dims()
+        spec = [i for i, s in enumerate(ta.t.shape) if isinstance(s, int)]
+
+        ta.specialize_dims(*spec)
+        ksel.return_tensor(ta.t).specialize_dims(*spec)
 
     def eager_execute(self, device_moniker, tensor):
         return tensor