diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 01b3c92..c12b9f3 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -35,7 +35,8 @@ jobs:
             find . -type f > before_test_files.txt
             conda env create -f environment.yml
             conda activate my_alignn
-            conda install pytest coverage codecov pandas numpy matplotlib phonopy scikit-learn jarvis-tools --quiet
+            conda install pytest coverage codecov pandas numpy matplotlib phonopy scikit-learn  --quiet
+            #conda install pytest coverage codecov pandas numpy matplotlib phonopy scikit-learn jarvis-tools --quiet
             export DGLBACKEND=pytorch
             export CUDA_VISIBLE_DEVICES="-1"
             pip install phonopy flake8 pytest pycodestyle pydocstyle codecov pytest-cov coverage
diff --git a/README.md b/README.md
index 0230fb4..0d07c5d 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,9 @@
-![alt text](https://github.com/usnistgov/alignn/actions/workflows/main.yml/badge.svg)
-[![codecov](https://codecov.io/gh/usnistgov/alignn/branch/main/graph/badge.svg?token=S5X4OYC80V)](https://codecov.io/gh/usnistgov/alignn)
+![alt text](https://github.com/atomgptlab/alignn/actions/workflows/main.yml/badge.svg)
+[![codecov](https://codecov.io/gh/atomgptlab/alignn/branch/main/graph/badge.svg?token=S5X4OYC80V)](https://codecov.io/gh/atomgptlab/alignn)
 [![PyPI version](https://badge.fury.io/py/alignn.svg)](https://badge.fury.io/py/alignn)
-![GitHub tag (latest by date)](https://img.shields.io/github/v/tag/usnistgov/alignn)
-![GitHub code size in bytes](https://img.shields.io/github/languages/code-size/usnistgov/alignn)
-![GitHub commit activity](https://img.shields.io/github/commit-activity/y/usnistgov/alignn)
+![GitHub tag (latest by date)](https://img.shields.io/github/v/tag/atomgptlab/alignn)
+![GitHub code size in bytes](https://img.shields.io/github/languages/code-size/atomgptlab/alignn)
+![GitHub commit activity](https://img.shields.io/github/commit-activity/y/atomgptlab/alignn)
 [![Downloads](https://pepy.tech/badge/alignn)](https://pepy.tech/project/alignn)
 <!--
 [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/atomistic-line-graph-neural-network-for/formation-energy-on-materials-project)](https://paperswithcode.com/sota/formation-energy-on-materials-project?p=atomistic-line-graph-neural-network-for)
@@ -35,7 +35,7 @@ Atomisitic line graph neural network-based FF (ALIGNN-FF) (https://pubs.rsc.org/
 
 
 
-![ALIGNN layer schematic](https://github.com/usnistgov/alignn/blob/develop/alignn/tex/schematic_lg.jpg)
+![ALIGNN layer schematic](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/schematic_lg.jpg)
 
 <a name="install"></a>
 Installation
@@ -72,7 +72,7 @@ You can laso install a development version of alignn by cloning the repository a
 conda create --name my_alignn python=3.10 -y
 conda activate my_alignn
 conda install dgl=2.1.0 pytorch torchvision torchaudio pytorch-cuda -c pytorch -c nvidia
-git clone https://github.com/usnistgov/alignn
+git clone https://github.com/atomgptlab/alignn
 cd alignn
 python -m pip install -e .
 ```
@@ -122,9 +122,9 @@ Users can keep their structure files in `POSCAR`, `.cif`, `.xyz` or `.pdb` files
 
 In this directory, `id_prop.csv`, the filenames, and correponding target values are kept in `comma separated values (csv) format`.
 
-Here is an example of training OptB88vdw bandgaps of 50 materials from JARVIS-DFT database. The example is created using the [generate_sample_data_reg.py](https://github.com/usnistgov/alignn/blob/main/alignn/examples/sample_data/scripts/generate_sample_data_reg.py) script. Users can modify the script for more than 50 data, or make their own dataset in this format. For list of available datasets see [Databases](https://jarvis-tools.readthedocs.io/en/master/databases.html).
+Here is an example of training OptB88vdw bandgaps of 50 materials from JARVIS-DFT database. The example is created using the [generate_sample_data_reg.py](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/scripts/generate_sample_data_reg.py) script. Users can modify the script for more than 50 data, or make their own dataset in this format. For list of available datasets see [Databases](https://jarvis-tools.readthedocs.io/en/master/databases.html).
 
-The dataset in split in 80:10:10 as training-validation-test set (controlled by `train_ratio, val_ratio, test_ratio`) . To change the split proportion and other parameters, change the [config_example.json](https://github.com/usnistgov/alignn/blob/main/alignn/examples/sample_data/config_example.json) file. If, users want to train on certain sets and val/test on another dataset, set `n_train`, `n_val`, `n_test` manually in the `config_example.json` and also set `keep_data_order` as True there so that random shuffle is disabled.
+The dataset in split in 80:10:10 as training-validation-test set (controlled by `train_ratio, val_ratio, test_ratio`) . To change the split proportion and other parameters, change the [config_example.json](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/config_example.json) file. If, users want to train on certain sets and val/test on another dataset, set `n_train`, `n_val`, `n_test` manually in the `config_example.json` and also set `keep_data_order` as True there so that random shuffle is disabled.
 
 A brief help guide (`-h`) can be obtained as follows.
 
@@ -169,7 +169,7 @@ To finetune model, use `--restart_model_path` tag as well in the above with the
 train_alignn.py --root_dir "alignn/examples/sample_data_ff" --restart_model_path "temp/best_model.pt" --config "alignn/examples/sample_data_ff/config_example_atomwise.json" --output_dir="temp1"
 ```
 
-Starting version v2024.10.30, we also allow global training for multi-output along with energy (graph wise output), forces (atomwise gradients), charges/magnetic moments etc. (atomwise but non-gradients) properties with or without additional fingerprints/features in graph. See examples [here](https://github.com/usnistgov/alignn/tree/main/alignn/examples). 
+Starting version v2024.10.30, we also allow global training for multi-output along with energy (graph wise output), forces (atomwise gradients), charges/magnetic moments etc. (atomwise but non-gradients) properties with or without additional fingerprints/features in graph. See examples [here](https://github.com/atomgptlab/alignn/tree/main/alignn/examples). 
 
 Multi-GPU training is allowed with `DistributedDataParallel` with `torchrun` command. This feature is not thoroughly tested yet. 
 Example:
@@ -179,7 +179,7 @@ torchrun --nproc_per_node=4 train_alignn.py --root_dir DataDir --config config.j
 ```
 For multi-GPU training make sure you have correct SLURM/PBS script setup correctly such as `#SBATCH -n 4, #SBATCH -N 1, #SBATCH --gres=gpu:4` etc.
 
-High-throughput like training: Users can also try training using multiple example scripts to run multiple dataset (such as JARVIS-DFT, Materials project, QM9_JCTC etc.). Look into the [alignn/scripts/train_*.py](https://github.com/usnistgov/alignn/tree/main/alignn/scripts) folder. This is done primarily to make the trainings more automated rather than making folder/ csv files etc.
+High-throughput like training: Users can also try training using multiple example scripts to run multiple dataset (such as JARVIS-DFT, Materials project, QM9_JCTC etc.). Look into the [alignn/scripts/train_*.py](https://github.com/atomgptlab/alignn/tree/main/alignn/scripts) folder. This is done primarily to make the trainings more automated rather than making folder/ csv files etc.
 These scripts automatically download datasets from [Databases in jarvis-tools](https://jarvis-tools.readthedocs.io/en/master/databases.html) and train several models. Make sure you specify your specific queuing system details in the scripts.
 
 
@@ -193,7 +193,7 @@ Using pre-trained models
 
 All the trained models are distributed on [Figshare](https://figshare.com/projects/ALIGNN_models/126478.
 
-The [pretrained.py script](https://github.com/usnistgov/alignn/blob/develop/alignn/pretrained.py) can be applied to use them. These models can be used to directly make predictions.
+The [pretrained.py script](https://github.com/atomgptlab/alignn/blob/develop/alignn/pretrained.py) can be applied to use them. These models can be used to directly make predictions.
 
 A brief help section (`-h`) is shown using:
 
@@ -226,18 +226,18 @@ Several supporting scripts for stucture optimization, equation of states, phonon
 Web-apps
 ------------
 
-A basic web-app is for direct-prediction available at [JARVIS-ALIGNN app](https://jarvis.nist.gov/jalignn/). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom and bandgap using data trained on JARVIS-DFT dataset.
+A basic web-app is for direct-prediction available at [AtomGPT.org ALIGNN app](https://atomgpt.org/alignn). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom, bandgap and other properties using data trained on JARVIS-DFT dataset.
 
-Similarly, a web-app for [ALIGNN-FF](https://jarvis.nist.gov/jalignnff/) for structure optimization is also available.
+Similarly, a web-app for [ALIGNN-FF](https://atomgpt.org/alignn_ff_dynamics) for structure optimization is also available.
 
-![JARVIS-ALIGNN](https://github.com/usnistgov/alignn/blob/develop/alignn/tex/jalignn.PNG)
+![JARVIS-ALIGNN](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/jalignn.PNG)
 
 
 
 <a name="alignnff"></a>
 ALIGNN-FF ASE Calculaor
 -------------------------
-![ALIGNN FF](https://github.com/usnistgov/alignn/blob/develop/alignn/tex/animation.gif)
+![ALIGNN FF](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/animation.gif)
 
 
 [ASE calculator](https://wiki.fysik.dtu.dk/ase/ase/calculators/calculators.html) provides interface to various codes. An example for ALIGNN-FF is give below. Note that there are multiple pretrained ALIGNN-FF models available, here we use the deafult_path model. As more accurate models are developed, they will be made available as well:
@@ -374,7 +374,7 @@ The results from models other than ALIGNN are reported as given in corresponding
 
 ### 4) On QM9 dataset
 
-Note the [issue](https://github.com/usnistgov/alignn/issues/54) related to QM9 dataset. The results from models other than ALIGNN are reported as given in corresponding papers, not necessarily reproduced by us. These models were trained with same parameters as solid-state databases but for 1000 epochs.
+Note the [issue](https://github.com/atomgptlab/alignn/issues/54) related to QM9 dataset. The results from models other than ALIGNN are reported as given in corresponding papers, not necessarily reproduced by us. These models were trained with same parameters as solid-state databases but for 1000 epochs.
 
 | Target | Units | SchNet | MEGNet  | DimeNet++ | ALIGNN |
 |:------:|-------|--------|---------|-----------|--------|
@@ -440,7 +440,7 @@ Useful notes (based on some of the queries we received)
 2) While comnventional '.cif' and '.pdb' files can be read using jarvis-tools, for complex files you might have to install `cif2cell` and `pytraj` respectively i.e.`pip install cif2cell==2.0.0a3` and `conda install -c ambermd pytraj`.
 3) Make sure you use `batch_size` as 32 or 64 for large datasets, and not 2 as given in the example config file, else it will take much longer to train, and performnce might drop a lot.
 4) Note that `train_alignn.py` and `pretrained.py` in alignn folder are actually python executable scripts. So, even if you don't provide absolute path of these scripts, they should work.
-5) Learn about the issue with QM9 results here: https://github.com/usnistgov/alignn/issues/54
+5) Learn about the issue with QM9 results here: https://github.com/atomgptlab/alignn/issues/54
 6) Make sure you have `pandas` version as >1.2.3.
 7) Starting March 2024, pytroch-ignite dependency will be removed to enable conda-forge build.
 
@@ -459,7 +459,7 @@ References
 8) [Unified graph neural network force-field for the periodic table](https://pubs.rsc.org/en/content/articlehtml/2023/dd/d2dd00096b)
 9) [Large Scale Benchmark of Materials Design Methods](https://www.nature.com/articles/s41524-024-01259-w)
 10) [Prediction of Magnetic Properties in van der Waals Magnets using Graph Neural Networks](https://doi.org/10.1103/PhysRevMaterials.8.114002)
-11) [CHIPS-FF: Benchmarking universal force-fields](https://github.com/usnistgov/chipsff)
+11) [CHIPS-FF: Benchmarking universal force-fields](https://github.com/atomgptlab/chipsff)
 
 
 Please see detailed publications list [here](https://jarvis-tools.readthedocs.io/en/master/publications.html).
@@ -468,13 +468,13 @@ Please see detailed publications list [here](https://jarvis-tools.readthedocs.io
 How to contribute
 -----------------
 
-For detailed instructions, please see [Contribution instructions](https://github.com/usnistgov/jarvis/blob/master/Contribution.rst)
+For detailed instructions, please see [Contribution instructions](https://github.com/atomgptlab/jarvis/blob/master/Contribution.rst)
 
 <a name="corres"></a>
 Correspondence
 --------------------
 
-Please report bugs as Github issues (https://github.com/usnistgov/alignn/issues) or email to kamal.choudhary@nist.gov.
+Please report bugs as Github issues (https://github.com/atomgptlab/alignn/issues) or email to drkamal@jhu.edu .
 
 <a name="fund"></a>
 Funding support
@@ -487,4 +487,4 @@ NIST-CHIPS (https://www.nist.gov/chips)
 Code of conduct
 --------------------
 
-Please see [Code of conduct](https://github.com/usnistgov/jarvis/blob/master/CODE_OF_CONDUCT.md)
+Please see [Code of conduct](https://github.com/atomgptlab/jarvis/blob/master/CODE_OF_CONDUCT.md)
diff --git a/alignn/__init__.py b/alignn/__init__.py
index ca6a785..0c35900 100644
--- a/alignn/__init__.py
+++ b/alignn/__init__.py
@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2025.4.1"
+__version__ = "2026.4.1"
diff --git a/alignn/ff/all_models_alignn.json b/alignn/ff/all_models_alignn.json
index b145a5f..3c7c872 100644
--- a/alignn/ff/all_models_alignn.json
+++ b/alignn/ff/all_models_alignn.json
@@ -1,59 +1,59 @@
 {
-    "jv_formation_energy_peratom_alignn": "https://figshare.com/ndownloader/files/31458679",
-    "jv_optb88vdw_total_energy_alignn": "https://figshare.com/ndownloader/files/31459642",
-    "jv_optb88vdw_bandgap_alignn": "https://figshare.com/ndownloader/files/31459636",
-    "jv_mbj_bandgap_alignn": "https://figshare.com/ndownloader/files/31458694",
-    "jv_spillage_alignn": "https://figshare.com/ndownloader/files/31458736",
-    "jv_slme_alignn": "https://figshare.com/ndownloader/files/31458727",
-    "jv_bulk_modulus_kv_alignn": "https://figshare.com/ndownloader/files/31458649",
-    "jv_shear_modulus_gv_alignn": "https://figshare.com/ndownloader/files/31458724",
-    "jv_n-Seebeck_alignn": "https://figshare.com/ndownloader/files/31458718",
-    "jv_n-powerfact_alignn": "https://figshare.com/ndownloader/files/31458712",
-    "intermat_cbm": "https://figshare.com/ndownloader/files/45392908",
-    "intermat_vbm": "https://figshare.com/ndownloader/files/45392914",
-    "intermat_phi": "https://figshare.com/ndownloader/files/45392911",
-    "jv_magmom_oszicar_alignn": "https://figshare.com/ndownloader/files/31458685",
-    "jv_kpoint_length_unit_alignn": "https://figshare.com/ndownloader/files/31458682",
-    "jv_avg_elec_mass_alignn": "https://figshare.com/ndownloader/files/31458643",
-    "jv_avg_hole_mass_alignn": "https://figshare.com/ndownloader/files/31458646",
-    "jv_epsx_alignn": "https://figshare.com/ndownloader/files/31458667",
-    "jv_mepsx_alignn": "https://figshare.com/ndownloader/files/31458703",
-    "jv_max_efg_alignn": "https://figshare.com/ndownloader/files/31458691",
-    "jv_ehull_alignn": "https://figshare.com/ndownloader/files/31458658",
-    "jv_dfpt_piezo_max_dielectric_alignn": "https://figshare.com/ndownloader/files/31458652",
-    "jv_dfpt_piezo_max_dij_alignn": "https://figshare.com/ndownloader/files/31458655",
-    "jv_exfoliation_energy_alignn": "https://figshare.com/ndownloader/files/31458676",
-    "jv_supercon_tc_alignn": "https://figshare.com/ndownloader/files/38789199",
-    "jv_supercon_edos_alignn": "https://figshare.com/ndownloader/files/39946300",
-    "jv_supercon_debye_alignn": "https://figshare.com/ndownloader/files/39946297",
-    "jv_supercon_a2F_alignn": "https://figshare.com/ndownloader/files/38801886",
-    "mp_e_form_alignn": "https://figshare.com/ndownloader/files/31458811",
-    "mp_gappbe_alignn": "https://figshare.com/ndownloader/files/31458814",
-    "tinnet_O_alignn": "https://figshare.com/ndownloader/files/41962800",
-    "tinnet_N_alignn": "https://figshare.com/ndownloader/files/41962797",
-    "tinnet_OH_alignn": "https://figshare.com/ndownloader/files/41962803",
-    "AGRA_O_alignn": "https://figshare.com/ndownloader/files/41966619",
-    "AGRA_OH_alignn": "https://figshare.com/ndownloader/files/41966610",
-    "AGRA_CHO_alignn": "https://figshare.com/ndownloader/files/41966643",
-    "AGRA_CO_alignn": "https://figshare.com/ndownloader/files/41966634",
-    "AGRA_COOH_alignn": "https://figshare.com/ndownloader/41966646",
-    "qm9_U0_alignn": "https://figshare.com/ndownloader/files/31459054",
-    "qm9_U_alignn": "https://figshare.com/ndownloader/files/31459051",
-    "qm9_alpha_alignn": "https://figshare.com/ndownloader/files/31459027",
-    "qm9_gap_alignn": "https://figshare.com/ndownloader/files/31459036",
-    "qm9_G_alignn": "https://figshare.com/ndownloader/files/31459033",
-    "qm9_HOMO_alignn": "https://figshare.com/ndownloader/files/31459042",
-    "qm9_LUMO_alignn": "https://figshare.com/ndownloader/files/31459045",
-    "qm9_ZPVE_alignn": "https://figshare.com/ndownloader/files/31459057",
-    "hmof_co2_absp_alignn": "https://figshare.com/ndownloader/files/31459198",
-    "hmof_max_co2_adsp_alignn": "https://figshare.com/ndownloader/files/31459207",
-    "hmof_surface_area_m2g_alignn": "https://figshare.com/ndownloader/files/31459222",
-    "hmof_surface_area_m2cm3_alignn": "https://figshare.com/ndownloader/files/31459219",
-    "hmof_pld_alignn": "https://figshare.com/ndownloader/files/31459216",
-    "hmof_lcd_alignn": "https://figshare.com/ndownloader/files/31459201",
-    "hmof_void_fraction_alignn": "https://figshare.com/ndownloader/files/31459228",
-    "ocp2020_all": "https://figshare.com/ndownloader/files/41411025",
-    "ocp2020_100k": "https://figshare.com/ndownloader/files/41967303",
-    "ocp2020_10k": "https://figshare.com/ndownloader/files/41967330",
-    "jv_pdos_alignn": "https://figshare.com/ndownloader/files/36757005"
-}
\ No newline at end of file
+    "jv_formation_energy_peratom_alignn": "https://ndownloader.figshare.com/files/31458679",
+    "jv_optb88vdw_total_energy_alignn": "https://ndownloader.figshare.com/files/31459642",
+    "jv_optb88vdw_bandgap_alignn": "https://ndownloader.figshare.com/files/31459636",
+    "jv_mbj_bandgap_alignn": "https://ndownloader.figshare.com/files/31458694",
+    "jv_spillage_alignn": "https://ndownloader.figshare.com/files/31458736",
+    "jv_slme_alignn": "https://ndownloader.figshare.com/files/31458727",
+    "jv_bulk_modulus_kv_alignn": "https://ndownloader.figshare.com/files/31458649",
+    "jv_shear_modulus_gv_alignn": "https://ndownloader.figshare.com/files/31458724",
+    "jv_n-Seebeck_alignn": "https://ndownloader.figshare.com/files/31458718",
+    "jv_n-powerfact_alignn": "https://ndownloader.figshare.com/files/31458712",
+    "intermat_cbm": "https://ndownloader.figshare.com/files/45392908",
+    "intermat_vbm": "https://ndownloader.figshare.com/files/45392914",
+    "intermat_phi": "https://ndownloader.figshare.com/files/45392911",
+    "jv_magmom_oszicar_alignn": "https://ndownloader.figshare.com/files/31458685",
+    "jv_kpoint_length_unit_alignn": "https://ndownloader.figshare.com/files/31458682",
+    "jv_avg_elec_mass_alignn": "https://ndownloader.figshare.com/files/31458643",
+    "jv_avg_hole_mass_alignn": "https://ndownloader.figshare.com/files/31458646",
+    "jv_epsx_alignn": "https://ndownloader.figshare.com/files/31458667",
+    "jv_mepsx_alignn": "https://ndownloader.figshare.com/files/31458703",
+    "jv_max_efg_alignn": "https://ndownloader.figshare.com/files/31458691",
+    "jv_ehull_alignn": "https://ndownloader.figshare.com/files/31458658",
+    "jv_dfpt_piezo_max_dielectric_alignn": "https://ndownloader.figshare.com/files/31458652",
+    "jv_dfpt_piezo_max_dij_alignn": "https://ndownloader.figshare.com/files/31458655",
+    "jv_exfoliation_energy_alignn": "https://ndownloader.figshare.com/files/31458676",
+    "jv_supercon_tc_alignn": "https://ndownloader.figshare.com/files/38789199",
+    "jv_supercon_edos_alignn": "https://ndownloader.figshare.com/files/39946300",
+    "jv_supercon_debye_alignn": "https://ndownloader.figshare.com/files/39946297",
+    "jv_supercon_a2F_alignn": "https://ndownloader.figshare.com/files/38801886",
+    "mp_e_form_alignn": "https://ndownloader.figshare.com/files/31458811",
+    "mp_gappbe_alignn": "https://ndownloader.figshare.com/files/31458814",
+    "tinnet_O_alignn": "https://ndownloader.figshare.com/files/41962800",
+    "tinnet_N_alignn": "https://ndownloader.figshare.com/files/41962797",
+    "tinnet_OH_alignn": "https://ndownloader.figshare.com/files/41962803",
+    "AGRA_O_alignn": "https://ndownloader.figshare.com/files/41966619",
+    "AGRA_OH_alignn": "https://ndownloader.figshare.com/files/41966610",
+    "AGRA_CHO_alignn": "https://ndownloader.figshare.com/files/41966643",
+    "AGRA_CO_alignn": "https://ndownloader.figshare.com/files/41966634",
+    "AGRA_COOH_alignn": "https://ndownloader.figshare.com/41966646",
+    "qm9_U0_alignn": "https://ndownloader.figshare.com/files/31459054",
+    "qm9_U_alignn": "https://ndownloader.figshare.com/files/31459051",
+    "qm9_alpha_alignn": "https://ndownloader.figshare.com/files/31459027",
+    "qm9_gap_alignn": "https://ndownloader.figshare.com/files/31459036",
+    "qm9_G_alignn": "https://ndownloader.figshare.com/files/31459033",
+    "qm9_HOMO_alignn": "https://ndownloader.figshare.com/files/31459042",
+    "qm9_LUMO_alignn": "https://ndownloader.figshare.com/files/31459045",
+    "qm9_ZPVE_alignn": "https://ndownloader.figshare.com/files/31459057",
+    "hmof_co2_absp_alignn": "https://ndownloader.figshare.com/files/31459198",
+    "hmof_max_co2_adsp_alignn": "https://ndownloader.figshare.com/files/31459207",
+    "hmof_surface_area_m2g_alignn": "https://ndownloader.figshare.com/files/31459222",
+    "hmof_surface_area_m2cm3_alignn": "https://ndownloader.figshare.com/files/31459219",
+    "hmof_pld_alignn": "https://ndownloader.figshare.com/files/31459216",
+    "hmof_lcd_alignn": "https://ndownloader.figshare.com/files/31459201",
+    "hmof_void_fraction_alignn": "https://ndownloader.figshare.com/files/31459228",
+    "ocp2020_all": "https://ndownloader.figshare.com/files/41411025",
+    "ocp2020_100k": "https://ndownloader.figshare.com/files/41967303",
+    "ocp2020_10k": "https://ndownloader.figshare.com/files/41967330",
+    "jv_pdos_alignn": "https://ndownloader.figshare.com/files/36757005"
+}
diff --git a/alignn/ff/all_models_alignn_atomwise.json b/alignn/ff/all_models_alignn_atomwise.json
index 7d42fb8..54a0165 100644
--- a/alignn/ff/all_models_alignn_atomwise.json
+++ b/alignn/ff/all_models_alignn_atomwise.json
@@ -1,28 +1,28 @@
 {
-    "v12.2.2024_dft_3d_307k": "https://figshare.com/ndownloader/files/50904240",
-    "v12.2.2024_mp_1.5mill": "https://figshare.com/ndownloader/files/50904783",
-    "v12.2.2024_mp_187k": "https://figshare.com/ndownloader/files/50904801",
-    "v2024.12.12_dft_3d_multi_prop": "https://figshare.com/ndownloader/files/52025186",
-    "alex_band_gap": "https://figshare.com/ndownloader/files/51993641",
-    "alex_dos_pa": "https://figshare.com/ndownloader/files/51993653",
-    "alex_e_form": "https://figshare.com/ndownloader/files/51993647",
-    "alex_e_hull": "https://figshare.com/ndownloader/files/51993644",
-    "alex_e_total": "https://figshare.com/ndownloader/files/51993650",
-    "alex_mag_per_vol": "https://figshare.com/ndownloader/files/51993659",
-    "alex_vol_pa": "https://figshare.com/ndownloader/files/51993656",
-    "v10.30.2024_dft_3d_307k": "https://figshare.com/ndownloader/files/50634327",
-    "v10.30.2024_mp_168k": "https://figshare.com/ndownloader/files/50634318",
-    "v8.29.2024_dft_3d": "https://figshare.com/ndownloader/files/48889834",
-    "v8.29.2024_mpf": "https://figshare.com/ndownloader/files/48889837",
-    "v5.27.2024": "https://figshare.com/ndownloader/files/47286127",
-    "alignnff_fmult": "https://figshare.com/ndownloader/files/41583585",
-    "alignnff_wt10": "https://figshare.com/ndownloader/files/41583594",
-    "alignnff_fd": "https://figshare.com/ndownloader/files/41583582",
-    "alignnff_wt01": "https://figshare.com/ndownloader/files/41583588",
-    "alignnff_wt1": "https://figshare.com/ndownloader/files/41583591",
-    "fmult_mlearn_only": "https://figshare.com/ndownloader/files/41583597",
-    "aff_Oct23": "https://figshare.com/ndownloader/files/42880573",
-    "revised": "https://figshare.com/ndownloader/files/41583600",
-    "scf_fd_top_10_en_42_fmax_600_wt01": "https://figshare.com/ndownloader/files/41967375",
-    "scf_fd_top_10_en_42_fmax_600_wt10": "https://figshare.com/ndownloader/files/41967372"
+    "v12.2.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50904240",
+    "v12.2.2024_mp_1.5mill": "https://ndownloader.figshare.com/files/50904783",
+    "v12.2.2024_mp_187k": "https://ndownloader.figshare.com/files/50904801",
+    "v2024.12.12_dft_3d_multi_prop": "https://ndownloader.figshare.com/files/52025186",
+    "alex_band_gap": "https://ndownloader.figshare.com/files/51993641",
+    "alex_dos_pa": "https://ndownloader.figshare.com/files/51993653",
+    "alex_e_form": "https://ndownloader.figshare.com/files/51993647",
+    "alex_e_hull": "https://ndownloader.figshare.com/files/51993644",
+    "alex_e_total": "https://ndownloader.figshare.com/files/51993650",
+    "alex_mag_per_vol": "https://ndownloader.figshare.com/files/51993659",
+    "alex_vol_pa": "https://ndownloader.figshare.com/files/51993656",
+    "v10.30.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50634327",
+    "v10.30.2024_mp_168k": "https://ndownloader.figshare.com/files/50634318",
+    "v8.29.2024_dft_3d": "https://ndownloader.figshare.com/files/48889834",
+    "v8.29.2024_mpf": "https://ndownloader.figshare.com/files/48889837",
+    "v5.27.2024": "https://ndownloader.figshare.com/files/47286127",
+    "alignnff_fmult": "https://ndownloader.figshare.com/files/41583585",
+    "alignnff_wt10": "https://ndownloader.figshare.com/files/41583594",
+    "alignnff_fd": "https://ndownloader.figshare.com/files/41583582",
+    "alignnff_wt01": "https://ndownloader.figshare.com/files/41583588",
+    "alignnff_wt1": "https://ndownloader.figshare.com/files/41583591",
+    "fmult_mlearn_only": "https://ndownloader.figshare.com/files/41583597",
+    "aff_Oct23": "https://ndownloader.figshare.com/files/42880573",
+    "revised": "https://ndownloader.figshare.com/files/41583600",
+    "scf_fd_top_10_en_42_fmax_600_wt01": "https://ndownloader.figshare.com/files/41967375",
+    "scf_fd_top_10_en_42_fmax_600_wt10": "https://ndownloader.figshare.com/files/41967372"
 }
diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index 85f77ff..ebc1e6c 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -19,6 +19,7 @@
 import numpy as np
 from tqdm import tqdm
 import torch
+from jarvis.core.utils import get_cache_dir
 
 # Reference: https://doi.org/10.1039/D2DD00096B
 
@@ -44,7 +45,8 @@ def get_figshare_model_ff(
     all_models_ff = get_all_models()
     # https://doi.org/10.6084/m9.figshare.23695695
     if dir_path is None:
-        dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        # dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        dir_path = os.path.join(get_cache_dir("alignn_ff"), model_name)
     # cwd=os.getcwd()
     dir_path = os.path.abspath(dir_path)
     if not os.path.exists(dir_path):
@@ -262,6 +264,7 @@ def __init__(
                     torch.load(
                         os.path.join(path, model_filename),
                         map_location=self.device,
+                        weights_only=False,
                     )
                 )
             else:
@@ -269,9 +272,10 @@ def __init__(
                     torch.load(
                         os.path.join(path, model_filename),
                         map_location=self.device,
+                        weights_only=False,
                     )["model"]
                 )
-            model.to(device)
+            model.to(self.device)
             model.eval()
             self.model = model
         else:
@@ -289,6 +293,10 @@ def calculate(self, atoms, properties=None, system_changes=None):
             atom_features=self.config["atom_features"],
             use_canonize=self.config["use_canonize"],
         )
+        # print("self.devicee", self.device)
+        # print("g", g.device)
+        # print("lg", lg.device)
+        # print("model", self.model)
 
         if self.config["model"]["alignn_layers"] > 0:
             result = self.model(
@@ -445,9 +453,11 @@ def __init__(
             torch.load(
                 os.path.join(ff_path, ff_model_filename),
                 map_location=self.device,
+                weights_only=False,
             )
         )
         ff_model.eval()
+        ff_model.to(self.device)
         self.ff_model = ff_model
         if prop_path is None and prop_model is None:
             prop_path = get_figshare_model_ff(
@@ -466,9 +476,11 @@ def __init__(
             torch.load(
                 os.path.join(prop_path, prop_model_filename),
                 map_location=self.device,
+                weights_only=False,
             )
         )
         prop_model.eval()
+        prop_model.to(self.device)
         self.prop_model = prop_model
 
     def calculate(
diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index d18eef2..b065e4f 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -16,7 +16,9 @@
 from ase.optimize.gpmin.gpmin import GPMin
 from ase.optimize.lbfgs import LBFGS, LBFGSLineSearch
 from ase.optimize.mdmin import MDMin
-from ase.constraints import ExpCellFilter
+
+# from ase.constraints import ExpCellFilter
+from ase.filters import ExpCellFilter
 from ase.eos import EquationOfState
 from ase.units import kJ
 from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG
@@ -45,7 +47,7 @@
 from matplotlib.gridspec import GridSpec
 from sklearn.metrics import mean_absolute_error
 from tqdm import tqdm
-
+from jarvis.core.utils import get_cache_dir
 # import torch
 from alignn.ff.calculators import (
     AlignnAtomwiseCalculator,
@@ -83,7 +85,8 @@ def get_figshare_model_ff(
     all_models_ff = get_all_models()
     # https://doi.org/10.6084/m9.figshare.23695695
     if dir_path is None:
-        dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        # dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        dir_path = os.path.join(get_cache_dir("alignn_ff"), model_name)
     # cwd=os.getcwd()
     dir_path = os.path.abspath(dir_path)
     if not os.path.exists(dir_path):
@@ -143,7 +146,8 @@ def get_figshare_model_prop(
     all_models_prop = get_all_models_prop()
     # https://doi.org/10.6084/m9.figshare.23695695
     if dir_path is None:
-        dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        # dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        dir_path = os.path.join(get_cache_dir("alignn_models"), model_name)
     # cwd=os.getcwd()
     dir_path = os.path.abspath(dir_path)
     if not os.path.exists(dir_path):
@@ -456,12 +460,24 @@ def run_nvt_langevin(
             self.set_momentum_maxwell_boltzmann(
                 temperature_K=initial_temperature_K
             )
+        # self.dyn = Langevin(
+        #    self.atoms,
+        #    self.timestep,
+        #    temperature_K=temperature_K,
+        #    friction=friction,
+        #    comm=self.communicator,
+        #    # communicator=self.communicator,
+        # )
+        langevin_kwargs = dict(
+            temperature_K=temperature_K,
+            friction=friction,
+        )
+        if self.communicator is not None:
+            langevin_kwargs["comm"] = self.communicator
         self.dyn = Langevin(
             self.atoms,
             self.timestep,
-            temperature_K=temperature_K,
-            friction=friction,
-            communicator=self.communicator,
+            **langevin_kwargs,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -489,12 +505,25 @@ def run_nvt_andersen(
             self.set_momentum_maxwell_boltzmann(
                 temperature_K=initial_temperature_K
             )
+        # self.dyn = Andersen(
+        #    self.atoms,
+        #    self.timestep,
+        #    temperature_K=temperature_K,
+        #    andersen_prob=andersen_prob,
+        #    comm=self.communicator,
+        #    # communicator=self.communicator,
+        # )
+
+        andersen_kwargs = dict(
+            temperature_K=temperature_K,
+            andersen_prob=andersen_prob,
+        )
+        if self.communicator is not None:
+            andersen_kwargs["comm"] = self.communicator
         self.dyn = Andersen(
             self.atoms,
             self.timestep,
-            temperature_K=temperature_K,
-            andersen_prob=andersen_prob,
-            communicator=self.communicator,
+            **andersen_kwargs,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -529,7 +558,8 @@ def run_nvt_berendsen(
             self.timestep,
             temperature_K=temperature_K,
             taut=taut,
-            communicator=self.communicator,
+            comm=self.communicator,
+            # communicator=self.communicator,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -560,15 +590,31 @@ def run_npt_berendsen(
             self.set_momentum_maxwell_boltzmann(
                 temperature_K=initial_temperature_K
             )
-        self.dyn = NPTBerendsen(
-            self.atoms,
-            self.timestep,
+        # self.dyn = NPTBerendsen(
+        #    self.atoms,
+        #    self.timestep,
+        #    temperature_K=temperature_K,
+        #    taut=taut,
+        #    taup=taup,
+        #    pressure=pressure,
+        #    compressibility=compressibility,
+        #    comm=self.communicator,
+        #    # communicator=self.communicator,
+        # )
+
+        npt_kwargs = dict(
             temperature_K=temperature_K,
             taut=taut,
             taup=taup,
             pressure=pressure,
             compressibility=compressibility,
-            communicator=self.communicator,
+        )
+        if self.communicator is not None:
+            npt_kwargs["comm"] = self.communicator
+        self.dyn = NPTBerendsen(
+            self.atoms,
+            self.timestep,
+            **npt_kwargs,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index bbee7c1..aff0243 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -343,9 +343,10 @@ def __init__(
                 config.hidden_features, config.additional_output_features
             )
         if self.classification:
-            self.fc = nn.Linear(config.hidden_features, 1)
-            self.softmax = nn.Sigmoid()
-            # self.softmax = nn.LogSoftmax(dim=1)
+            # self.fc = nn.Linear(config.hidden_features, 1)
+            self.fc = nn.Linear(config.hidden_features, 2)
+            # self.softmax = nn.Sigmoid()
+            self.softmax = nn.LogSoftmax(dim=1)
         else:
             self.fc = nn.Linear(config.hidden_features, config.output_features)
         self.link = None
@@ -650,6 +651,8 @@ def forward(
 
         if self.classification:
             # out = torch.max(out,dim=1)
+            if out.dim() == 1:
+                out = out.unsqueeze(0)
             out = self.softmax(out)
         result["out"] = out
         result["additional"] = additional_out
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index d016e2c..c05a9d5 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -14,11 +14,15 @@
 import argparse
 from jarvis.core.atoms import Atoms
 from alignn.graphs import Graph
-from jarvis.db.jsonutils import dumpjson
+from jarvis.db.jsonutils import loadjson, dumpjson
 import pandas as pd
 from alignn.dataset import get_torch_dataset
 import numpy as np
-
+from alignn.models.alignn_atomwise import (
+    ALIGNNAtomWise,
+    ALIGNNAtomWiseConfig,
+)
+from jarvis.core.utils import get_cache_dir
 # from jarvis.core.graphs import Graph
 
 tqdm.pandas()
@@ -35,159 +39,163 @@
 # to load as a calculator
 all_models = {
     "jv_formation_energy_peratom_alignn": [
-        "https://figshare.com/ndownloader/files/31458679",
+        "https://ndownloader.figshare.com/files/31458679",
         1,
     ],
     "jv_optb88vdw_total_energy_alignn": [
-        "https://figshare.com/ndownloader/files/31459642",
+        "https://ndownloader.figshare.com/files/31459642",
         1,
     ],
     "jv_optb88vdw_bandgap_alignn": [
-        "https://figshare.com/ndownloader/files/31459636",
+        "https://ndownloader.figshare.com/files/31459636",
         1,
     ],
     "jv_mbj_bandgap_alignn": [
-        "https://figshare.com/ndownloader/files/31458694",
+        "https://ndownloader.figshare.com/files/31458694",
         1,
     ],
     "jv_spillage_alignn": [
-        "https://figshare.com/ndownloader/files/31458736",
+        "https://ndownloader.figshare.com/files/31458736",
         1,
     ],
-    "jv_slme_alignn": ["https://figshare.com/ndownloader/files/31458727", 1],
+    "jv_slme_alignn": ["https://ndownloader.figshare.com/files/31458727", 1],
     "jv_bulk_modulus_kv_alignn": [
-        "https://figshare.com/ndownloader/files/31458649",
+        "https://ndownloader.figshare.com/files/31458649",
         1,
     ],
     "jv_shear_modulus_gv_alignn": [
-        "https://figshare.com/ndownloader/files/31458724",
+        "https://ndownloader.figshare.com/files/31458724",
         1,
     ],
     "jv_n-Seebeck_alignn": [
-        "https://figshare.com/ndownloader/files/31458718",
+        "https://ndownloader.figshare.com/files/31458718",
         1,
     ],
     "jv_n-powerfact_alignn": [
-        "https://figshare.com/ndownloader/files/31458712",
+        "https://ndownloader.figshare.com/files/31458712",
         1,
     ],
     "intermat_cbm": [
-        "https://figshare.com/ndownloader/files/45392908",
+        "https://ndownloader.figshare.com/files/45392908",
         1,
     ],
     "intermat_vbm": [
-        "https://figshare.com/ndownloader/files/45392914",
+        "https://ndownloader.figshare.com/files/45392914",
         1,
     ],
     "intermat_phi": [
-        "https://figshare.com/ndownloader/files/45392911",
+        "https://ndownloader.figshare.com/files/45392911",
         1,
     ],
     "jv_magmom_oszicar_alignn": [
-        "https://figshare.com/ndownloader/files/31458685",
+        "https://ndownloader.figshare.com/files/31458685",
         1,
     ],
     "jv_kpoint_length_unit_alignn": [
-        "https://figshare.com/ndownloader/files/31458682",
+        "https://ndownloader.figshare.com/files/31458682",
         1,
     ],
     "jv_avg_elec_mass_alignn": [
-        "https://figshare.com/ndownloader/files/31458643",
+        "https://ndownloader.figshare.com/files/31458643",
         1,
     ],
     "jv_avg_hole_mass_alignn": [
-        "https://figshare.com/ndownloader/files/31458646",
+        "https://ndownloader.figshare.com/files/31458646",
         1,
     ],
-    "jv_epsx_alignn": ["https://figshare.com/ndownloader/files/31458667", 1],
-    "jv_mepsx_alignn": ["https://figshare.com/ndownloader/files/31458703", 1],
+    "jv_epsx_alignn": ["https://ndownloader.figshare.com/files/31458667", 1],
+    "jv_mepsx_alignn": ["https://ndownloader.figshare.com/files/31458703", 1],
     "jv_max_efg_alignn": [
-        "https://figshare.com/ndownloader/files/31458691",
+        "https://ndownloader.figshare.com/files/31458691",
         1,
     ],
-    "jv_ehull_alignn": ["https://figshare.com/ndownloader/files/31458658", 1],
+    "jv_ehull_alignn": ["https://ndownloader.figshare.com/files/31458658", 1],
     "jv_dfpt_piezo_max_dielectric_alignn": [
-        "https://figshare.com/ndownloader/files/31458652",
+        "https://ndownloader.figshare.com/files/31458652",
         1,
     ],
     "jv_dfpt_piezo_max_dij_alignn": [
-        "https://figshare.com/ndownloader/files/31458655",
+        "https://ndownloader.figshare.com/files/31458655",
         1,
     ],
     "jv_exfoliation_energy_alignn": [
-        "https://figshare.com/ndownloader/files/31458676",
+        "https://ndownloader.figshare.com/files/31458676",
         1,
     ],
     "jv_supercon_tc_alignn": [
-        "https://figshare.com/ndownloader/files/38789199",
+        "https://ndownloader.figshare.com/files/38789199",
         1,
     ],
     "jv_supercon_edos_alignn": [
-        "https://figshare.com/ndownloader/files/39946300",
+        "https://ndownloader.figshare.com/files/39946300",
         1,
     ],
     "jv_supercon_debye_alignn": [
-        "https://figshare.com/ndownloader/files/39946297",
+        "https://ndownloader.figshare.com/files/39946297",
         1,
     ],
     "jv_supercon_a2F_alignn": [
-        "https://figshare.com/ndownloader/files/38801886",
+        "https://ndownloader.figshare.com/files/38801886",
         100,
     ],
     "mp_e_form_alignn": [
-        "https://figshare.com/ndownloader/files/31458811",
+        "https://ndownloader.figshare.com/files/31458811",
         1,
     ],
     "mp_gappbe_alignn": [
-        "https://figshare.com/ndownloader/files/31458814",
+        "https://ndownloader.figshare.com/files/31458814",
         1,
     ],
-    "tinnet_O_alignn": ["https://figshare.com/ndownloader/files/41962800", 1],
-    "tinnet_N_alignn": ["https://figshare.com/ndownloader/files/41962797", 1],
-    "tinnet_OH_alignn": ["https://figshare.com/ndownloader/files/41962803", 1],
-    "AGRA_O_alignn": ["https://figshare.com/ndownloader/files/41966619", 1],
-    "AGRA_OH_alignn": ["https://figshare.com/ndownloader/files/41966610", 1],
-    "AGRA_CHO_alignn": ["https://figshare.com/ndownloader/files/41966643", 1],
-    "AGRA_CO_alignn": ["https://figshare.com/ndownloader/files/41966634", 1],
-    "AGRA_COOH_alignn": ["https://figshare.com/ndownloader/41966646", 1],
-    "qm9_U0_alignn": ["https://figshare.com/ndownloader/files/31459054", 1],
-    "qm9_U_alignn": ["https://figshare.com/ndownloader/files/31459051", 1],
-    "qm9_alpha_alignn": ["https://figshare.com/ndownloader/files/31459027", 1],
-    "qm9_gap_alignn": ["https://figshare.com/ndownloader/files/31459036", 1],
-    "qm9_G_alignn": ["https://figshare.com/ndownloader/files/31459033", 1],
-    "qm9_HOMO_alignn": ["https://figshare.com/ndownloader/files/31459042", 1],
-    "qm9_LUMO_alignn": ["https://figshare.com/ndownloader/files/31459045", 1],
-    "qm9_ZPVE_alignn": ["https://figshare.com/ndownloader/files/31459057", 1],
+    "tinnet_O_alignn": ["https://ndownloader.figshare.com/files/41962800", 1],
+    "tinnet_N_alignn": ["https://ndownloader.figshare.com/files/41962797", 1],
+    "tinnet_OH_alignn": ["https://ndownloader.figshare.com/files/41962803", 1],
+    "AGRA_O_alignn": ["https://ndownloader.figshare.com/files/41966619", 1],
+    "AGRA_OH_alignn": ["https://ndownloader.figshare.com/files/41966610", 1],
+    "AGRA_CHO_alignn": ["https://ndownloader.figshare.com/files/41966643", 1],
+    "AGRA_CO_alignn": ["https://ndownloader.figshare.com/files/41966634", 1],
+    "AGRA_COOH_alignn": ["https://ndownloader.figshare.com/41966646", 1],
+    "qm9_U0_alignn": ["https://ndownloader.figshare.com/files/31459054", 1],
+    "qm9_U_alignn": ["https://ndownloader.figshare.com/files/31459051", 1],
+    "qm9_alpha_alignn": ["https://ndownloader.figshare.com/files/31459027", 1],
+    "qm9_gap_alignn": ["https://ndownloader.figshare.com/files/31459036", 1],
+    "qm9_G_alignn": ["https://ndownloader.figshare.com/files/31459033", 1],
+    "qm9_HOMO_alignn": ["https://ndownloader.figshare.com/files/31459042", 1],
+    "qm9_LUMO_alignn": ["https://ndownloader.figshare.com/files/31459045", 1],
+    "qm9_ZPVE_alignn": ["https://ndownloader.figshare.com/files/31459057", 1],
     "hmof_co2_absp_alignn": [
-        "https://figshare.com/ndownloader/files/31459198",
+        "https://ndownloader.figshare.com/files/31459198",
         5,
     ],
     "hmof_max_co2_adsp_alignn": [
-        "https://figshare.com/ndownloader/files/31459207",
+        "https://ndownloader.figshare.com/files/31459207",
         1,
     ],
     "hmof_surface_area_m2g_alignn": [
-        "https://figshare.com/ndownloader/files/31459222",
+        "https://ndownloader.figshare.com/files/31459222",
         1,
     ],
     "hmof_surface_area_m2cm3_alignn": [
-        "https://figshare.com/ndownloader/files/31459219",
+        "https://ndownloader.figshare.com/files/31459219",
         1,
     ],
-    "hmof_pld_alignn": ["https://figshare.com/ndownloader/files/31459216", 1],
-    "hmof_lcd_alignn": ["https://figshare.com/ndownloader/files/31459201", 1],
+    "hmof_pld_alignn": ["https://ndownloader.figshare.com/files/31459216", 1],
+    "hmof_lcd_alignn": ["https://ndownloader.figshare.com/files/31459201", 1],
     "hmof_void_fraction_alignn": [
-        "https://figshare.com/ndownloader/files/31459228",
+        "https://ndownloader.figshare.com/files/31459228",
         1,
     ],
-    "ocp2020_all": ["https://figshare.com/ndownloader/files/41411025", 1],
-    "ocp2020_100k": ["https://figshare.com/ndownloader/files/41967303", 1],
-    "ocp2020_10k": ["https://figshare.com/ndownloader/files/41967330", 1],
+    "ocp2020_all": ["https://ndownloader.figshare.com/files/41411025", 1],
+    "ocp2020_100k": ["https://ndownloader.figshare.com/files/41967303", 1],
+    "ocp2020_10k": ["https://ndownloader.figshare.com/files/41967330", 1],
     "jv_pdos_alignn": [
-        "https://figshare.com/ndownloader/files/36757005",
+        "https://ndownloader.figshare.com/files/36757005",
         66,
         {"alignn_layers": 6, "gcn_layers": 6},
     ],
+    "jv_raman_alignn": [
+        "https://ndownloader.figshare.com/files/62805487",
+        200,
+    ],
 }
 
 
@@ -252,7 +260,8 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     # else:
     #    config_params = {}
     zfile = model_name + ".zip"
-    path = str(os.path.join(os.path.dirname(__file__), zfile))
+    # path = str(os.path.join(os.path.dirname(__file__), zfile))
+    path = os.path.join(get_cache_dir("alignn_models"), zfile)
     if not os.path.isfile(path):
         response = requests.get(url, stream=True)
         total_size_in_bytes = int(response.headers.get("content-length", 0))
@@ -283,6 +292,7 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     print("Path", os.path.abspath(path))
     print("Config", os.path.abspath(cfg))
     config = json.loads(zipfile.ZipFile(path).read(cfg))
+    print("config", config, type(config))
     # print("Loading the zipfile...", zipfile.ZipFile(path).namelist())
     data = zipfile.ZipFile(path).read(tmp)
     # model = ALIGNN(
@@ -290,16 +300,34 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     #        name="alignn", output_features=output_features, **config_params
     #    )
     # )
-    model = ALIGNN(ALIGNNConfig(**config["model"]))
-
-    new_file, filename = tempfile.mkstemp()
-    with open(filename, "wb") as f:
-        f.write(data)
-    model.load_state_dict(torch.load(filename, map_location=device)["model"])
-    model.to(device)
-    model.eval()
-    if os.path.exists(filename):
-        os.remove(filename)
+    print(config, type(config))
+    if config["model"]["name"] == "alignn":
+        model = ALIGNN(ALIGNNConfig(**config["model"]))
+        new_file, filename = tempfile.mkstemp()
+        with open(filename, "wb") as f:
+            f.write(data)
+        model.load_state_dict(
+            torch.load(filename, map_location=device, weights_only=False)[
+                "model"
+            ]
+        )
+        model.to(device)
+        model.eval()
+        if os.path.exists(filename):
+            os.remove(filename)
+    if config["model"]["name"] == "alignn_atomwise":
+        model = ALIGNNAtomWise(ALIGNNAtomWiseConfig(**config["model"]))
+        new_file, filename = tempfile.mkstemp()
+        with open(filename, "wb") as f:
+            f.write(data)
+        model.load_state_dict(
+            torch.load(filename, map_location=device, weights_only=False)
+        )
+        model.to(device)
+        model.eval()
+        if os.path.exists(filename):
+            os.remove(filename)
+
     return model
 
 
@@ -310,7 +338,21 @@ def get_prediction(
     max_neighbors=12,
 ):
     """Get model prediction on a single structure."""
-    model = get_figshare_model(model_name)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if os.path.isdir(model_name):
+
+        # import torch
+        # from jarvis.db.jsonutils import loadjson
+
+        config = loadjson(os.path.join(model_name, "config.json"))
+        model_path = os.path.join(model_name, "best_model.pt")
+        tmp = ALIGNNAtomWiseConfig(**config["model"])
+        model = ALIGNNAtomWise(tmp)
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model = model.to(device)
+        print(model)
+    else:
+        model = get_figshare_model(model_name)
     # print("Loading completed.")
     g, lg = Graph.atom_dgl_multigraph(
         atoms,
@@ -318,14 +360,11 @@ def get_prediction(
         max_neighbors=max_neighbors,
     )
     lat = torch.tensor(atoms.lattice_mat)
-    out_data = (
-        model([g.to(device), lg.to(device), lat.to(device)])
-        .detach()
-        .cpu()
-        .numpy()
-        .flatten()
-        .tolist()
-    )
+    out_data = model([g.to(device), lg.to(device), lat.to(device)])
+    if isinstance(out_data, dict):
+        out_data = out_data["out"]
+    print("out_data", out_data)
+    out_data = out_data.detach().cpu().numpy().flatten().tolist()
     return out_data
 
 
diff --git a/environment.yml b/environment.yml
index 886addb..f3f686f 100644
--- a/environment.yml
+++ b/environment.yml
@@ -60,7 +60,6 @@ dependencies:
   - idna=3.6=pyhd8ed1ab_0
   - importlib-metadata=7.1.0=pyha770c72_0
   - iniconfig=2.0.0=pyhd8ed1ab_0
-  - jarvis-tools=2024.3.24=pyhd8ed1ab_0
   - jinja2=3.1.3=pyhd8ed1ab_0
   - joblib=1.3.2=pyhd8ed1ab_0
   - jsonpatch=1.33=pyhd8ed1ab_0
@@ -229,7 +228,7 @@ dependencies:
   - pip:
       - annotated-types==0.6.0
       - appdirs==1.4.4
-      - ase==3.22.1
+      - ase
       - autograd==1.6.2
       - bokeh==2.4.2
       - cif2cell==2.1.0
@@ -240,6 +239,7 @@ dependencies:
       - flake8==7.0.0
       - fsspec==2024.3.1
       - future==1.0.0
+      - jarvis-tools
       - lightgbm==3.3.1
       - mccabe==0.7.0
       - mkdocs-material==9.5.14
diff --git a/setup.py b/setup.py
index b9c616c..ab9bb21 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 """ALIGNN: Atomistic LIne Graph Neural Network.
 
-https://jarvis.nist.gov.
+https://atomgpt.org/alignn
 """
 
 import setuptools
@@ -10,18 +10,20 @@
 
 setuptools.setup(
     name="alignn",
-    version="2025.4.1",
+    version="2026.4.1",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",
     install_requires=[
-        "numpy<2.0",
+        # "numpy<2.0",
         # "numpy>=1.19.5,<2.0.0",
+        "numpy>=1.19.5",
         "scipy>=1.6.1",
         "jarvis-tools>=2021.07.19",
-        "torch<=2.2.1",
+        "torch>=2.2.1",
+        # "torch<=2.2.1",
         "mpmath<=1.3.0",
-        "dgl<=1.1.1",
+        # "dgl>=0.6.0",
         "spglib>=2.0.2",
         "scikit-learn>=0.22.2",
         "matplotlib>=3.4.1",
@@ -32,7 +34,8 @@
         "flake8>=3.9.1",
         "pycodestyle>=2.7.0",
         "pydocstyle>=6.0.0",
-        "pyparsing>=2.2.1,<3",
+        # "pyparsing>=2.2.1,<3",
+        "pyparsing>=2.2.1",
         "ase",
         "lmdb",
     ],
@@ -43,7 +46,7 @@
     ],
     long_description=long_description,
     long_description_content_type="text/markdown",
-    url="https://github.com/usnistgov/alignn",
+    url="https://github.com/atomgptlab/alignn",
     packages=setuptools.find_packages(),
     package_data={
         "alignn.ff": [