How to utilize GPU

[keenanbar]

Is there some kind of guide or tutorial on utilizing the GPU in pytorch to perform some of the computation? I have looked at various tutorials and examples which mention the GPU and this one (https://docs.gempy.org/examples/real/Moureze.html#sphx-glr-download-examples-real-moureze-py) which actually provides the performance differences, but I have not been able to get gempy working with my GPU.

I would greatly appreciate it if anybody could share a minimum reproducible example of computing a model on the GPU, along with the dependencies and versions that work.

I have verified that torch is working properly because I can create some arrays and do a matrix multiply with them and see it computing on the GPU. At this point I have tried several different versions of pytorch and cuda but have not been able to get gempy to use the GPU. These are the statements I have included to try and get it working:

import torch
print(torch.cuda.is_available())  # this prints True

from gempy_engine.config import AvailableBackends

gp.compute_model(
    gempy_model=geo_model,
    engine_config=gp.data.GemPyEngineConfig(
        use_gpu=True,
        dtype='float32',
        backend=AvailableBackends.PYTORCH
    )
)

Edit: I am using a windows machine with an RTX3090. I create a conda virtual environment for gempy and its dependencies.

[NilsChudalla]

Hi @keenanbar,

thanks for your question. Can you show me the versions of packages you have installed? I would like to reproduce your issue... or maybe, we can already spot some problem with packages installed.

Kind regards
Nils

[keenanbar]

Here is the output from conda list. I think I grabbed all the correct packages. I think I grabbed all the packages I need, but it's possible I missed a critical one. I installed pykeops, but when I went to test it with pykeops.test_torch_bindings() I just got the usual issue on windows with fnctl. I'm not sure if pykeops is required for windows since I only saw it referenced under the Linux section of the pypi page for gempy-engine.

# Name                    Version                   Build  Channel
annotated-types           0.7.0                    pypi_0    pypi
bzip2                     1.0.8                h2bbff1b_6
ca-certificates           2026.3.19            haa95532_0
certifi                   2022.12.7                pypi_0    pypi
charset-normalizer        2.1.1                    pypi_0    pypi
contourpy                 1.3.3                    pypi_0    pypi
cycler                    0.12.1                   pypi_0    pypi
filelock                  3.25.2                   pypi_0    pypi
fonttools                 4.62.1                   pypi_0    pypi
fsspec                    2026.2.0                 pypi_0    pypi
gempy                     2025.2.0                 pypi_0    pypi
gempy-engine              2025.2.0                 pypi_0    pypi
gempy-viewer              2025.1.6                 pypi_0    pypi
idna                      3.4                      pypi_0    pypi
jinja2                    3.1.6                    pypi_0    pypi
keopscore                 2.3                      pypi_0    pypi
kiwisolver                1.5.0                    pypi_0    pypi
libexpat                  2.7.5                hd7fb8db_0
libffi                    3.4.4                hd77b12b_1
libzlib                   1.3.1                h02ab6af_0
markupsafe                3.0.3                    pypi_0    pypi
matplotlib                3.10.8                   pypi_0    pypi
mpmath                    1.3.0                    pypi_0    pypi
networkx                  3.6.1                    pypi_0    pypi
numpy                     1.26.0                   pypi_0    pypi
openssl                   3.5.6                hbb43b14_0
packaging                 26.0            py311haa95532_0
pandas                    3.0.2                    pypi_0    pypi
pillow                    12.1.1                   pypi_0    pypi
pip                       26.0.1             pyhc872135_1
pybind11                  3.0.3                    pypi_0    pypi
pydantic                  2.13.1                   pypi_0    pypi
pydantic-core             2.46.1                   pypi_0    pypi
pykeops                   2.3                      pypi_0    pypi
pyparsing                 3.3.2                    pypi_0    pypi
python                    3.11.15              h1044e36_0
python-dateutil           2.9.0.post0              pypi_0    pypi
python-dotenv             1.2.2                    pypi_0    pypi
requests                  2.28.1                   pypi_0    pypi
setuptools                82.0.1          py311haa95532_0
six                       1.17.0                   pypi_0    pypi
sqlite                    3.51.2               hee5a0db_0
sympy                     1.14.0                   pypi_0    pypi
tk                        8.6.15               hf199647_0
torch                     2.1.0+cu121              pypi_0    pypi
torchaudio                2.1.0+cu121              pypi_0    pypi
torchvision               0.16.0+cu121             pypi_0    pypi
typing-extensions         4.15.0                   pypi_0    pypi
typing-inspection         0.4.2                    pypi_0    pypi
tzdata                    2026.1                   pypi_0    pypi
ucrt                      10.0.22621.0         haa95532_0
urllib3                   1.26.13                  pypi_0    pypi
vc                        14.3                h2df5915_10
vc14_runtime              14.44.35208         h4927774_10
vs2015_runtime            14.44.35208         ha6b5a95_10
wheel                     0.46.3          py311haa95532_0
xz                        5.8.2                h53af0af_0
zlib                      1.3.1                h02ab6af_0

[javoha]

Small Addon question: What kind of error do you get when executing the code you provided?

[keenanbar]

I should have put that in the original post. I don't actually get any error at all. It computes the model just fine. However, all of the computation is being done on the CPU. When I watch the GPU the memory usage and utilization don't change at all during execution.

[keenanbar]

model2_moreorientations_gempycompat.json
@NilsChudalla @javoha
I included a small MRE so that you can take a look and see if you observe the same behavior that I do. I also included the JSON that you can use to run it. When I run this, I don't get any errors, but it seems everything is happening on the CPU. The GPU memory and utilization don't change at all when the script is running. I attached a list of all the packages in the environment as well.

While I was troubleshooting I did run into something interesting. When I uncomment the torch.set_default_device('cuda') command, I get the following error.

line 45, in _compute_rest_ref_in_numpy
    ref_points = sp_coords[partitions_bool]
                 ~~~~~~~~~^^^^^^^^^^^^^^^^^
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

I'm not sure exactly how pytorch is integrated with gempy and if this is to be expected if some tensors would always be calculated on the CPU regardless of the use_gpu parameter, but I included it just in case that gives you a clue as to what may be going wrong.

MRE:

import torch
print(torch.cuda.is_available())  # this prints True
print(torch.cuda.get_device_name(0))  # this prints: NVIDIA GeForce RTX 3090 Ti
# torch.set_default_device('cuda')

import gempy as gp
import gempy_viewer as gpv
from gempy.modules.json_io.json_operations import JsonIO
from gempy_engine.config import AvailableBackends


# Set the JSON path
new_json_path = r'model2_moreorientations_gempycompat.json'

# Create the geo model from the new JSON file
geo_model = JsonIO.load_model_from_json(new_json_path)

# Compute the model
gp.compute_model(
    gempy_model=geo_model,
    engine_config=gp.data.GemPyEngineConfig(
        use_gpu=True,
        dtype='float32',
        backend=AvailableBackends.PYTORCH
    )
)

# Display the model
gpv.plot_3d(geo_model, show_surfaces=True, image=False, show_lith=False)

Environment:

# Name                    Version                   Build  Channel
annotated-types           0.7.0                    pypi_0    pypi
attrs                     26.1.0                   pypi_0    pypi
bzip2                     1.0.8                h2bbff1b_6  
ca-certificates           2026.3.19            haa95532_0  
certifi                   2022.12.7                pypi_0    pypi
charset-normalizer        2.1.1                    pypi_0    pypi
contourpy                 1.3.3                    pypi_0    pypi
cycler                    0.12.1                   pypi_0    pypi
cyclopts                  4.10.2                   pypi_0    pypi
docstring-parser          0.18.0                   pypi_0    pypi
docutils                  0.22.4                   pypi_0    pypi
filelock                  3.25.2                   pypi_0    pypi
fonttools                 4.62.1                   pypi_0    pypi
fsspec                    2026.2.0                 pypi_0    pypi
gempy                     2025.2.0                 pypi_0    pypi
gempy-engine              2025.2.0                 pypi_0    pypi
gempy-viewer              2025.1.6                 pypi_0    pypi
idna                      3.4                      pypi_0    pypi
jinja2                    3.1.6                    pypi_0    pypi
kiwisolver                1.5.0                    pypi_0    pypi
libexpat                  2.7.5                hd7fb8db_0  
libffi                    3.4.4                hd77b12b_1  
libzlib                   1.3.1                h02ab6af_0  
markdown-it-py            4.0.0                    pypi_0    pypi
markupsafe                3.0.3                    pypi_0    pypi
matplotlib                3.10.8                   pypi_0    pypi
mdurl                     0.1.2                    pypi_0    pypi
mpmath                    1.3.0                    pypi_0    pypi
networkx                  3.6.1                    pypi_0    pypi
numpy                     1.26.0                   pypi_0    pypi
openssl                   3.5.6                hbb43b14_0  
packaging                 26.0            py311haa95532_0  
pandas                    3.0.2                    pypi_0    pypi
pillow                    12.1.1                   pypi_0    pypi
pip                       26.0.1             pyhc872135_1  
platformdirs              4.9.6                    pypi_0    pypi
pooch                     1.9.0                    pypi_0    pypi
pybind11                  3.0.3                    pypi_0    pypi
pydantic                  2.13.1                   pypi_0    pypi
pydantic-core             2.46.1                   pypi_0    pypi
pygments                  2.20.0                   pypi_0    pypi
pyparsing                 3.3.2                    pypi_0    pypi
python                    3.11.15              h1044e36_0  
python-dateutil           2.9.0.post0              pypi_0    pypi
python-dotenv             1.2.2                    pypi_0    pypi
pyvista                   0.47.3                   pypi_0    pypi
requests                  2.28.1                   pypi_0    pypi
rich                      15.0.0                   pypi_0    pypi
rich-rst                  1.3.2                    pypi_0    pypi
scooby                    0.11.1                   pypi_0    pypi
setuptools                82.0.1          py311haa95532_0  
six                       1.17.0                   pypi_0    pypi
sqlite                    3.51.2               hee5a0db_0  
sympy                     1.14.0                   pypi_0    pypi
tk                        8.6.15               hf199647_0  
torch                     2.1.0+cu121              pypi_0    pypi
torchaudio                2.1.0+cu121              pypi_0    pypi
torchvision               0.16.0+cu121             pypi_0    pypi
typing-extensions         4.15.0                   pypi_0    pypi
typing-inspection         0.4.2                    pypi_0    pypi
tzdata                    2026.1                   pypi_0    pypi
ucrt                      10.0.22621.0         haa95532_0  
urllib3                   1.26.13                  pypi_0    pypi
vc                        14.3                h2df5915_10  
vc14_runtime              14.44.35208         h4927774_10  
vs2015_runtime            14.44.35208         ha6b5a95_10  
vtk                       9.6.1                    pypi_0    pypi
wheel                     0.46.3          py311haa95532_0  
xz                        5.8.2                h53af0af_0  
zlib                      1.3.1                h02ab6af_0  

[javoha]

Hi again,
first of all thanks a lot for looking into it in more detail. This is exactly the kind of contribution that helps with developing GemPy.

I actually have never tracked the usage but had significant speed-ups using setting GPU so I assumed it worked. You might be right that some things are always computed on CPU. This however is more of a backend question that falls into @Leguark expertise as the main developer. It might take some time for him or us to look into this however.

My suggestion for now would be if to check if you can track a computation speed increase for models with more points/higher resolution comparing cpu vs gpu turned on.

Sorry for the slightly unsatisfying answer,
Jan