one more

microsoft · Jan 17, 2025 · 5b86288 · 5b86288
1 parent 30bdcda
commit 5b86288
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/sampling_tests.cpp b/test/sampling_tests.cpp
@@ -461,7 +461,7 @@ TEST(SamplingTests, RandomizedSamplingTopKCuda) {
   // Run test
   for (int i = 0; i < num_iter; i++) {
     auto generator = Generators::CreateGenerator(*model, *params);
-    Generators::DeviceSpan<float> logits_gpu = params->p_device->Allocate<float>(config.model.vocab_size * batch_size);
+    Generators::DeviceSpan<float> logits_gpu = params->p_device->Allocate<float>(vocab_size * batch_size);
     auto cpu_span = logits_gpu.CpuSpan();
     // Shuffle integers 1 to k randomly into cpu_span
     for (int i = 0; i < batch_size; i++) {