Can anyone show me a script that training a m3gnet model using CUDA?

[SmallBearC]

I have tried as follows:

from future import annotations
import torch
torch.set_default_device('cuda')
import os
import shutil

import numpy as np
import pytorch_lightning as pl
from dgl.data.utils import split_dataset
from pymatgen.core import Structure
from matgl.ext.pymatgen import Structure2Graph, get_element_list
from matgl.graph.data import M3GNetDataset, MEGNetDataset, MGLDataLoader, collate_fn, collate_fn_efs
from matgl.models import M3GNet, MEGNet
from matgl.utils.training import ModelLightningModule, PotentialLightningModule

if name == 'main':
stru0 = Structure.from_file("../../strus/0/POSCAR")
stru1 = Structure.from_file("../../strus/1/POSCAR")
......

Or

from future import annotations
import torch
import os
import shutil

import numpy as np
import pytorch_lightning as pl
from dgl.data.utils import split_dataset
from pymatgen.core import Structure
from matgl.ext.pymatgen import Structure2Graph, get_element_list
from matgl.graph.data import M3GNetDataset, MEGNetDataset, MGLDataLoader, collate_fn, collate_fn_efs
from matgl.models import M3GNet, MEGNet
from matgl.utils.training import ModelLightningModule, PotentialLightningModule

if name == 'main':
torch.set_default_device('cuda')
torch.device('cuda')
stru0 = Structure.from_file("../../strus/0/POSCAR")
stru1 = Structure.from_file("../../strus/1/POSCAR")
......

But it all had the same error:

Original Traceback (most recent call last):
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
data = fetcher.fetch(index)
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 51, in
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/dgl/data/utils.py", line 425, in getitem
return self.dataset[self.indices[item]]
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/matgl/graph/data.py", line 321, in getitem
self.state_attr[idx],
RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

When I use the torch.device('cuda') as follows:

from future import annotations
import torch
torch.device('cuda')
import os
import shutil

import numpy as np
import pytorch_lightning as pl
from dgl.data.utils import split_dataset
from pymatgen.core import Structure
from matgl.ext.pymatgen import Structure2Graph, get_element_list
from matgl.graph.data import M3GNetDataset, MEGNetDataset, MGLDataLoader, collate_fn, collate_fn_efs
from matgl.models import M3GNet, MEGNet
from matgl.utils.training import ModelLightningModule, PotentialLightningModule

if name == 'main':
stru0 = Structure.from_file("../../strus/0/POSCAR")
stru1 = Structure.from_file("../../strus/1/POSCAR")
......

Or

from future import annotations
import torch
import os
import shutil

import numpy as np
import pytorch_lightning as pl
from dgl.data.utils import split_dataset
from pymatgen.core import Structure
from matgl.ext.pymatgen import Structure2Graph, get_element_list
from matgl.graph.data import M3GNetDataset, MEGNetDataset, MGLDataLoader, collate_fn, collate_fn_efs
from matgl.models import M3GNet, MEGNet
from matgl.utils.training import ModelLightningModule, PotentialLightningModule

if name == 'main':
torch.device('cuda')
stru0 = Structure.from_file("../../strus/0/POSCAR")
stru1 = Structure.from_file("../../strus/1/POSCAR")
......

I get this error:

File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/matgl/models/_m3gnet.py", line 224, in forward
g.edata["bond_vec"] = bond_vec
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/dgl/view.py", line 227, in setitem
self._graph._set_e_repr(self._etid, self._edges, {key: val})
File "/home/lycui/anaconda3/envs/mgl/lib/python3.9/site-packages/dgl/heterograph.py", line 4460, in _set_e_repr
raise DGLError(
dgl._ffi.base.DGLError: Cannot assign edge feature "bond_vec" on device cpu to a graph on device cuda:0. Call DGLGraph.to() to copy the graph to the same device.

**So I'm confused, I don't know what to do.

Here are some packages version:
dgl 1.1.0+cu117
cryptography 41.0.1
matgl 0.6.2
pytorch-lightning 2.0.4
torch 2.0.1
tqdm 4.65.0
pymatgen 2023.5.31**

[shyuep]

Given that we have no information on what machine you are using, there really isn't much else we can help you with here. Pls read the documentation on cuda in PyTorch.

[SmallBearC]

Thank you very much for responding once again.
Previously, I was using Google Colab and my own server, setting up everything from scratch in a completely new environment, which led to the errors I encountered.
Recently, I've successfully installed matgl on my own Ubuntu subsystem and managed to run it smoothly using the GPU. Indeed, I've made slight modifications to the source code, although I'm uncertain if this was the correct approach. I think you're right—it's plausible that the machine itself is also a factor.

I am profoundly impressed by your work and have been utilizing it continuously since MEGnet was introduced. I am eagerly looking forward to future updates. I also hope that I can rely on your help in the future.

Once again, thank you.