Back to shared_mem_size_bytes from kwargs

xsuite · Nov 10, 2023 · 1f7eb6e · 1f7eb6e
1 parent dfc322a
commit 1f7eb6e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py
@@ -671,7 +671,7 @@ def to_function_arg(self, arg, value):
     def num_args(self):
         return len(self.description.args)
 
-    def __call__(self, shared_mem_size_bytes=None, **kwargs):
+    def __call__(self, **kwargs):
         assert len(kwargs.keys()) == self.num_args
         arg_list = []
         for arg in self.description.args:
@@ -683,7 +683,9 @@ def __call__(self, shared_mem_size_bytes=None, **kwargs):
         else:
             n_threads = self.description.n_threads
 
-        if shared_mem_size_bytes is None:
+        if "shared_mem_size_bytes" in kwargs.keys():
+            shared_mem_size_bytes = kwargs["shared_mem_size_bytes"]
+        else:
             shared_mem_size_bytes = self.shared_mem_size_bytes
 
         grid_size = int(np.ceil(n_threads / self.block_size))