Update nvidia options

trustimaging · Jul 24, 2024 · bb20a84 · bb20a84
1 parent a821bfd
commit bb20a84
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stride/__init__.py b/stride/__init__.py
@@ -131,7 +131,7 @@ async def loop(worker, shot_id):
 
         if using_gpu:
             deviceid = devices[worker.indices[1] % num_gpus]
-            if platform == 'nvidia-acc':
+            if platform in ['nvidia-acc', 'nvidia-cuda']:
                 devito_args = _kwargs.get('devito_args', {})
                 devito_args['deviceid'] = deviceid
                 _kwargs['devito_args'] = devito_args
@@ -312,7 +312,7 @@ async def loop(worker, shot_id):
 
             if using_gpu:
                 deviceid = devices[worker.indices[1] % num_gpus]
-                if platform == 'nvidia-acc':
+                if platform in ['nvidia-acc', 'nvidia-cuda']:
                     devito_args = _kwargs.get('devito_args', {})
                     devito_args['deviceid'] = deviceid
                     _kwargs['devito_args'] = devito_args