Single GPU Support

Warning The use of GPU and mutli-GPU is under developpement and hasn’t been thoroughly tested yet. Proceed with caution !

Using pytorch as a backend, QMCTorch can leverage GPU cards available on your hardware. You of course must have the CUDA version of pytorch installed (https://pytorch.org/)

Let’s first import everything and create a molecule

import torch
from torch import optim
from qmctorch.scf import Molecule
from qmctorch.wavefunction import SlaterJastrow
from qmctorch.sampler import Metropolis
from qmctorch.utils import (plot_energy, plot_data)
mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', unit='bohr', redo_scf=True)

The use of GPU acceleration has been streamlined in QMCTorch, the only modification you need to do on your code is to specify cuda=True in the declaration of the wave function and sampler, this will automatically port all the necesaary tensors to the GPU and offload all the corresponding operation there.

if torch.cuda.is_available():
    wf = SlaterJastrow(mol, cuda=True)
    sampler = Metropolis(nwalkers=100, nstep=500, step_size=0.25,
                     nelec=wf.nelec, ndim=wf.ndim,
                     init=mol.domain('atomic'),
                     move={'type': 'all-elec', 'proba': 'normal'},
                     cuda=True)
else:
    print('CUDA not available, install torch with cuda support to proceed')