1 files changed, 3 insertions, 3 deletions

diff --git a/lib/test_vmalloc.c b/lib/test_vmalloc.c
index 83cdcaa82bf6..f832b095afba 100644
--- a/lib/test_vmalloc.c
+++ b/lib/test_vmalloc.c
@@ -383,14 +383,14 @@ static void shuffle_array(int *arr, int n)
 static int test_func(void *private)
 {
         struct test_driver *t = private;
-        cpumask_t newmask = CPU_MASK_NONE;
         int random_array[ARRAY_SIZE(test_case_array)];
         int index, i, j, ret;
         ktime_t kt;
         u64 delta;
 
-        cpumask_set_cpu(t->cpu, &newmask);
-        set_cpus_allowed_ptr(current, &newmask);
+        ret = set_cpus_allowed_ptr(current, cpumask_of(t->cpu));
+        if (ret < 0)
+                pr_err("Failed to set affinity to %d CPU\n", t->cpu);
 
         for (i = 0; i < ARRAY_SIZE(test_case_array); i++)
                 random_array[i] = i;