Replace ones with rand to workaround the crash from sdpa kernel

pytorch · Dec 17, 2024 · bedecd8 · bedecd8
1 parent 6cb6af9
commit bedecd8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/extension/benchmark/apple/Benchmark/Tests/GenericTests.mm b/extension/benchmark/apple/Benchmark/Tests/GenericTests.mm
@@ -78,7 +78,7 @@ @implementation GenericTests
 
           const auto sizes = tensor_meta->sizes();
           tensors.emplace_back(
-              ones({sizes.begin(), sizes.end()}, tensor_meta->scalar_type()));
+              rand({sizes.begin(), sizes.end()}, tensor_meta->scalar_type()));
           XCTAssertEqual(module->set_input(tensors.back(), index), Error::Ok);
         } break;
         default:

diff --git a/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm b/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm
@@ -87,7 +87,7 @@ @implementation LLaMATests
                               tokensPerSecondMetric.tokenCount = 0;
                               const auto status = runner->generate(
                                   "Once upon a time",
-                                  128,
+                                  50,
                                   [=](const std::string &token) {
                                     tokensPerSecondMetric.tokenCount++;
                                   },