fix for bad multi-device behavior

tenstorrent · Feb 24, 2025 · 61fdecc · 61fdecc
1 parent aa09a6f
commit 61fdecc
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/ttnn/cpp/ttnn/device_operation.hpp b/ttnn/cpp/ttnn/device_operation.hpp
@@ -432,8 +432,9 @@ typename device_operation_t::tensor_return_value_t launch_on_multi_device(
     std::vector<tensor_return_value_t> outputs;
     outputs.reserve(num_shards);
 
-    for (auto shard_index = 0; shard_index < num_shards; shard_index++) {
-        auto device = storage.get_buffer_for_device_id(shard_index)->device();
+    for (const auto &[shard_index, buffer] : storage.buffers ) {
+        auto device = buffer->device();
+        auto shard_tensor_args = get_shard_tensor_args<device_operation_t>(shard_index, device, tensor_args);
         auto shard_tensor_args = get_shard_tensor_args<device_operation_t>(shard_index, device, tensor_args);
         outputs.push_back(launch_on_single_device<device_operation_t>(cq_id, operation_attributes, shard_tensor_args));
     }