1 files changed, 4 insertions, 4 deletions

diff --git a/litmus/nvidia_info.c b/litmus/nvidia_info.c
index 0b39dcc84115..7883296a7a18 100644
--- a/litmus/nvidia_info.c
+++ b/litmus/nvidia_info.c
@@ -444,13 +444,13 @@ static int init_nv_device_reg(void)
         mb();
 
 
-        for(i = 0; i < NV_DEVICE_NUM; ++i) {
+        for(i = 0; i < num_online_gpus(); ++i) {
                 raw_spin_lock_init(&NV_DEVICE_REG[i].lock);
                 INIT_BINHEAP_HANDLE(&NV_DEVICE_REG[i].owners, gpu_owner_max_priority_order);
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
                 {
-                        int default_cpu = litmus->map_gpu_to_cpu(i, NV_DEVICE_NUM);
+                        int default_cpu = litmus->map_gpu_to_cpu(i);
 
                         NV_DEVICE_REG[i].callback.func = nvidia_klmirqd_cb;
                         NV_DEVICE_REG[i].callback.arg = (void*)(long long)(i);
@@ -478,7 +478,7 @@ static int shutdown_nv_device_reg(void)
                 int i;
                 nv_device_registry_t *reg;
 
-                for (i = 0; i < NV_DEVICE_NUM; ++i) {
+                for (i = 0; i < num_online_gpus(); ++i) {
 
                         TRACE("Shutting down GPU %d.\n", i);
 
@@ -829,7 +829,7 @@ int reg_nv_device(int reg_device_id, int reg_action, struct task_struct *t)
 {
         int ret;
 
-        if((reg_device_id < NV_DEVICE_NUM) && (reg_device_id >= 0))
+        if((reg_device_id < num_online_gpus()) && (reg_device_id >= 0))
         {
                 if(reg_action)
                         ret = __reg_nv_device(reg_device_id, t);