Quick-fix to platform

trustimaging · Sep 2, 2024 · 5ad3421 · 5ad3421
1 parent 6ce10d2
commit 5ad3421
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/stride/__init__.py b/stride/__init__.py
@@ -132,7 +132,7 @@ async def loop(worker, shot_id):
         if using_gpu:
             deviceid = devices[worker.indices[1] % num_gpus]
             if platform in ['nvidia-acc', 'nvidia-cuda']:
-                devito_args = _kwargs.get('devito_args', {})
+                devito_args = _kwargs.get('devito_args', {}).copy()
                 devito_args['deviceid'] = deviceid
                 _kwargs['devito_args'] = devito_args
             elif platform == 'gpu':
@@ -313,7 +313,7 @@ async def loop(worker, shot_id):
             if using_gpu:
                 deviceid = devices[worker.indices[1] % num_gpus]
                 if platform in ['nvidia-acc', 'nvidia-cuda']:
-                    devito_args = _kwargs.get('devito_args', {})
+                    devito_args = _kwargs.get('devito_args', {}).copy()
                     devito_args['deviceid'] = deviceid
                     _kwargs['devito_args'] = devito_args
                 elif platform == 'gpu':
@@ -399,8 +399,8 @@ async def loop(worker, shot_id):
 
                 if using_gpu:
                     deviceid = devices[worker.indices[1] % num_gpus]
-                    if platform == 'nvidia-acc':
-                        devito_args = _kwargs.get('devito_args', {})
+                    if platform in ['nvidia-acc', 'nvidia-cuda']:
+                        devito_args = _kwargs.get('devito_args', {}).copy()
                         devito_args['deviceid'] = deviceid
                         _kwargs['devito_args'] = devito_args
                     elif platform == 'gpu':