#0: Unary fix

tenstorrent · Nov 26, 2024 · 76234da · 76234da
1 parent 1171a25
commit 76234da
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_device_operation.cpp b/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_device_operation.cpp
@@ -170,6 +170,9 @@ spec_return_value_t UnaryDeviceOperation::compute_output_specs(
 
 tensor_return_value_t UnaryDeviceOperation::create_output_tensors(
     const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    if (tensor_args.preallocated_output.has_value()) {
+        return *tensor_args.preallocated_output;
+    }
     return create_device_tensor(compute_output_specs(args, tensor_args), tensor_args.input.device());
 }