Merge pull request #95 from FlorianDeconinck/fix/boilerplate_on_gpu

[QOL] Fix for boilerplate to work on GPU
NOAA-GFDL · Jan 2, 2025 · a8a2e73 · a8a2e73
2 parents 76f53c8 + f99914a
commit a8a2e73
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/ndsl/boilerplate.py b/ndsl/boilerplate.py
@@ -16,6 +16,7 @@
     TileCommunicator,
     TilePartitioner,
 )
+from ndsl.optional_imports import cupy as cp
 
 
 def _get_factories(
@@ -74,7 +75,9 @@ def _get_factories(
 
     grid_indexing = GridIndexing.from_sizer_and_communicator(sizer, comm)
     stencil_factory = StencilFactory(config=stencil_config, grid_indexing=grid_indexing)
-    quantity_factory = QuantityFactory(sizer, np)
+    quantity_factory = QuantityFactory(
+        sizer, cp if stencil_config.is_gpu_backend else np
+    )
 
     return stencil_factory, quantity_factory