RuntimeError: CUDA error: an illegal memory access was encountered

[MusannaGalib]

Hello,

I installed allegro using the following commands.

# install wandb
!pip install wandb

# install nequip
!pip install nequip==0.5.5 torch==1.11

# fix colab imports
import site
site.main()

# set to allow anonymous WandB
import os
os.environ["WANDB_ANONYMOUS"] = "must"

# install allegro
!git clone --depth 1 https://github.com/mir-group/allegro.git
!pip install allegro/

Allgero_model.txt

When I am training a model after loading the dataset it shows a CUDA memory error though I trained similar size model before. Any suggestions on this issue? Thanks

!CUDA_LAUNCH_BLOCKING=1 TORCH_CUDA_ALLOC_SYNC=1 nequip-train allegro/configs/tutorial.yaml

wandb: wandb version 0.16.3 is available!  To upgrade, please run:
wandb:  $ pip install wandb --upgrade
wandb: Tracking run with wandb version 0.15.12
Torch device: cuda
Processing dataset...
Loaded data: Batch(atomic_numbers=[7584316, 1], batch=[7584316], cell=[96004, 3, 3], edge_cell_shift=[95306240, 3], edge_index=[2, 95306240], forces=[7584316, 3], pbc=[96004, 3], pos=[7584316, 3], ptr=[96005], total_energy=[96004, 1], virial=[96004, 3, 3])
    processed data size: ~2842.24 MB
Cached processed data to disk
Done!
Successfully loaded the data set of type ASEDataset(96004)...
Replace string dataset_forces_rms to 0.7151474952697754
Replace string dataset_per_atom_total_energy_mean to -7.149214267730713
Atomic outputs are scaled by: [O, Al: 0.715147], shifted by [O, Al: -7.149214].
/home/galibubc/ML/lib/python3.10/site-packages/nequip/nn/_grad_output.py:196: UserWarning: !! Stresses in NequIP are in BETA and UNDER DEVELOPMENT: _please_ carefully check the sanity of your results and report any (potential) issues on the GitHub
  warnings.warn(
Replace string dataset_forces_rms to 0.7151474952697754
Initially outputs are globally scaled by: 0.7151474952697754, total_energy are globally shifted by None.
Successfully built the network...
Number of weights: 37416
Number of trainable weights: 37416
Traceback (most recent call last):
  File "/home/galibubc/ML/bin/nequip-train", line 8, in <module>
    sys.exit(main())
  File "/home/galibubc/ML/lib/python3.10/site-packages/nequip/scripts/train.py", line 78, in main
    trainer.train()
  File "/home/galibubc/ML/lib/python3.10/site-packages/nequip/train/trainer.py", line 759, in train
    self.init()
  File "/home/galibubc/ML/lib/python3.10/site-packages/nequip/train/trainer_wandb.py", line 14, in init
    super().init()
  File "/home/galibubc/ML/lib/python3.10/site-packages/nequip/train/trainer.py", line 719, in init
    self.init_objects()
  File "/home/galibubc/ML/lib/python3.10/site-packages/nequip/train/trainer.py", line 431, in init_objects
    self.ema = ExponentialMovingAverage(
  File "/home/galibubc/ML/lib/python3.10/site-packages/torch_ema/ema.py", line 52, in __init__
    self.shadow_params = [
  File "/home/galibubc/ML/lib/python3.10/site-packages/torch_ema/ema.py", line 53, in <listcomp>
    p.clone().detach()
RuntimeError: CUDA error: an illegal memory access was encountered
wandb: Waiting for W&B process to finish... (failed 1). Press Control-C to abort syncing.
wandb: 🚀 View run alumina at: https://wandb.ai/anony-moose-203399138991916818/Alumina_potential/runs/bvfmcm2q?apiKey=591b70bcf101936f96f9d405f0fd58342ab1bcc6
wandb: Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
wandb: Find logs at: ./wandb/run-20240227_173401-bvfmcm2q/logs

print(torch.__version__)
print(torch.version.cuda)
print(allegro.__version__)
print(nequip.__version__)
1.11.0+cu113
11.3
0.2.0
0.5.6

[MusannaGalib]

If someone faces a similar problem, the solution is - Reinstalling the virtual environment and packages solved the problem.