diff --git a/atomgpt/examples/forward_model/config.json b/atomgpt/examples/forward_model/config.json
index 446d714..b938a94 100644
--- a/atomgpt/examples/forward_model/config.json
+++ b/atomgpt/examples/forward_model/config.json
@@ -2,18 +2,18 @@
     "id_prop_path": "atomgpt/examples/forward_model/id_prop.csv",
     "prefix": "atomgpt_run",
     "model_name": "gpt2",
-    "batch_size": 5,
+    "batch_size": 2,
     "max_length": 512,
-    "num_epochs": 3,
+    "num_epochs": 1,
     "latent_dim": 1024,
     "learning_rate": 0.001,
     "test_each_run": true,
     "include_struct": false,
     "pretrained_path": "",
     "seed_val": 42,
-    "n_train": 40,
-    "n_val": 5,
-    "n_test": 5,
+    "n_train": 2,
+    "n_val": 2,
+    "n_test": 2,
     "output_dir": "atomgpt/examples/forward_model/out",
     "train_ratio": null,
     "val_ratio": 0.1,
diff --git a/atomgpt/inverse_models/kernels/__init__.py b/atomgpt/inverse_models/kernels/__init__.py
index cb04377..fb49219 100644
--- a/atomgpt/inverse_models/kernels/__init__.py
+++ b/atomgpt/inverse_models/kernels/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .cross_entropy_loss import fast_cross_entropy_loss
+from atomgpt.inverse_models.kernels.cross_entropy_loss import fast_cross_entropy_loss
 from .rms_layernorm import fast_rms_layernorm
 from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
 from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
diff --git a/environment.yml b/environment.yml
index 05e9ec1..3eb5829 100644
--- a/environment.yml
+++ b/environment.yml
@@ -1,365 +1,196 @@
 name: atomgpt
 channels:
+  - xformers
   - pytorch
   - nvidia
   - conda-forge
   - defaults
 dependencies:
   - _libgcc_mutex=0.1=conda_forge
-  - _openmp_mutex=4.5=2_kmp_llvm
-  - absl-py=2.1.0=pyhd8ed1ab_0
-  - aiohttp=3.9.5=py310h2372a71_0
-  - aiosignal=1.3.1=pyhd8ed1ab_0
-  - alignn=2024.4.20=pyhd8ed1ab_0
-  - annotated-types=0.7.0=pyhd8ed1ab_0
-  - aom=3.9.1=hac33072_0
-  - ase=3.23.0=pyhd8ed1ab_0
-  - astunparse=1.6.3=pyhd8ed1ab_0
-  - async-timeout=4.0.3=pyhd8ed1ab_0
-  - attrs=23.2.0=pyh71513ae_0
-  - aws-c-auth=0.7.22=h9137712_5
-  - aws-c-cal=0.6.15=h88a6e22_0
-  - aws-c-common=0.9.19=h4ab18f5_0
-  - aws-c-compression=0.2.18=h83b837d_6
-  - aws-c-event-stream=0.4.2=h0cbf018_13
-  - aws-c-http=0.8.2=h360477d_2
-  - aws-c-io=0.14.9=h2d549f9_2
-  - aws-c-mqtt=0.10.4=hf85b563_6
-  - aws-c-s3=0.5.10=h679ed35_3
-  - aws-c-sdkutils=0.1.16=h83b837d_2
-  - aws-checksums=0.1.18=h83b837d_6
-  - aws-crt-cpp=0.26.12=h8bc9c4d_0
-  - aws-sdk-cpp=1.11.329=hf74b5d1_5
-  - azure-core-cpp=1.12.0=h830ed8b_0
-  - azure-identity-cpp=1.8.0=hdb0d106_1
-  - azure-storage-blobs-cpp=12.11.0=ha67cba7_1
-  - azure-storage-common-cpp=12.6.0=he3f277c_1
-  - azure-storage-files-datalake-cpp=12.10.0=h29b5301_1
-  - babel=2.14.0=pyhd8ed1ab_0
-  - blinker=1.8.2=pyhd8ed1ab_0
-  - brotli=1.1.0=hd590300_1
-  - brotli-bin=1.1.0=hd590300_1
-  - brotli-python=1.1.0=py310hc6cd4ac_1
-  - bzip2=1.0.8=hd590300_5
-  - c-ares=1.28.1=hd590300_0
-  - ca-certificates=2024.6.2=hbcca054_0
-  - cached-property=1.5.2=hd8ed1ab_1
-  - cached_property=1.5.2=pyha770c72_1
-  - cairo=1.18.0=hbb29018_2
-  - certifi=2024.6.2=pyhd8ed1ab_0
-  - charset-normalizer=3.3.2=pyhd8ed1ab_0
-  - click=8.1.7=unix_pyh707e725_0
-  - colorama=0.4.6=pyhd8ed1ab_0
-  - contourpy=1.2.1=py310hd41b1e2_0
-  - cuda-cudart=12.4.127=0
-  - cuda-cupti=12.4.127=0
-  - cuda-libraries=12.4.0=0
-  - cuda-nvrtc=12.4.127=0
-  - cuda-nvtx=12.4.127=0
-  - cuda-opencl=12.5.39=0
-  - cuda-runtime=12.4.0=0
-  - cuda-version=12.5=3
-  - cycler=0.12.1=pyhd8ed1ab_0
-  - dav1d=1.2.1=hd590300_0
-  - dgl=2.1.0=py310h039f8d5_1
-  - dill=0.3.7=pyhd8ed1ab_0
-  - dm-tree=0.1.8=py310ha8c1f0e_4
-  - expat=2.6.2=h59595ed_0
-  - ffmpeg=7.0.1=gpl_h3589c0a_102
+  - _openmp_mutex=4.5=2_gnu
+  - blas=1.0=mkl
+  - bzip2=1.0.8=h7f98852_4
+  - ca-certificates=2024.2.2=hbcca054_0
+  - cairo=1.18.0=h3faef2a_0
+  - cffi=1.16.0=py39h7a31438_0
+  - cuda-cudart=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.4.99=0
+  - cuda-runtime=12.1.0=0
+  - cudatoolkit=11.7.0=hd8887f6_10
+  - expat=2.5.0=hcb278e6_1
   - filelock=3.15.4=pyhd8ed1ab_0
-  - flake8=7.1.0=pyhd8ed1ab_0
-  - flask=3.0.3=pyhd8ed1ab_0
-  - flatbuffers=24.3.25=h59595ed_0
   - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
   - font-ttf-inconsolata=3.000=h77eed37_0
   - font-ttf-source-code-pro=2.038=h77eed37_0
-  - font-ttf-ubuntu=0.83=h77eed37_2
+  - font-ttf-ubuntu=0.83=hab24e00_0
   - fontconfig=2.14.2=h14ed4e7_0
   - fonts-conda-ecosystem=1=0
   - fonts-conda-forge=1=0
-  - fonttools=4.53.0=py310hc51659f_0
   - freetype=2.12.1=h267a509_2
-  - fribidi=1.0.10=h36c2ea0_0
-  - frozenlist=1.4.1=py310h2372a71_0
-  - gast=0.5.4=pyhd8ed1ab_0
-  - gettext=0.22.5=h59595ed_2
-  - gettext-tools=0.22.5=h59595ed_2
-  - gflags=2.2.2=he1b5a44_1004
-  - ghp-import=2.1.0=pyhd8ed1ab_0
-  - giflib=5.2.2=hd590300_0
-  - glog=0.7.1=hbabe93e_0
-  - gmp=6.3.0=hac33072_2
-  - gmpy2=2.1.5=py310hc7909c9_1
-  - gnutls=3.7.9=hb077bed_0
-  - google-pasta=0.2.0=pyh8c360ce_0
-  - graphite2=1.3.13=h59595ed_1003
-  - grpcio=1.62.2=py310h1b8f574_0
-  - h5py=3.11.0=nompi_py310hf054cd7_102
-  - harfbuzz=8.5.0=hfac3d4d_0
-  - hdf5=1.14.3=nompi_hdf9ad27_105
-  - huggingface_hub=0.23.4=pyhd8ed1ab_0
+  - gettext=0.21.1=h27087fc_0
+  - gmp=6.3.0=h59595ed_1
+  - gmpy2=2.1.2=py39h376b7d2_1
   - icu=73.2=h59595ed_0
-  - idna=3.7=pyhd8ed1ab_0
-  - importlib-metadata=7.2.0=pyha770c72_0
-  - importlib-resources=6.4.0=pyhd8ed1ab_0
-  - importlib_resources=6.4.0=pyhd8ed1ab_0
-  - itsdangerous=2.2.0=pyhd8ed1ab_0
-  - jarvis-tools=2024.4.20=pyhd8ed1ab_0
+  - intel-openmp=2022.1.0=h9e868ea_3769
   - jinja2=3.1.4=pyhd8ed1ab_0
-  - joblib=1.4.2=pyhd8ed1ab_0
-  - keras=3.1.0=pyhd8ed1ab_0
-  - keyutils=1.6.1=h166bdaf_0
-  - kiwisolver=1.4.5=py310hd41b1e2_1
-  - krb5=1.21.2=h659d440_0
-  - lame=3.100=h166bdaf_1003
-  - lcms2=2.16=hb7c19ff_0
-  - ld_impl_linux-64=2.40=hf3520f5_7
-  - lerc=4.0.0=h27087fc_0
-  - libabseil=20240116.2=cxx17_h59595ed_0
-  - libaec=1.1.3=h59595ed_0
-  - libarrow=16.1.0=h4a673ee_10_cpu
-  - libarrow-acero=16.1.0=hac33072_10_cpu
-  - libarrow-dataset=16.1.0=hac33072_10_cpu
-  - libarrow-substrait=16.1.0=h7e0c224_10_cpu
-  - libasprintf=0.22.5=h661eb56_2
-  - libasprintf-devel=0.22.5=h661eb56_2
-  - libass=0.17.1=h8fe9dca_1
-  - libblas=3.9.0=22_linux64_openblas
-  - libbrotlicommon=1.1.0=hd590300_1
-  - libbrotlidec=1.1.0=hd590300_1
-  - libbrotlienc=1.1.0=hd590300_1
-  - libcblas=3.9.0=22_linux64_openblas
-  - libcrc32c=1.1.2=h9c3ff4c_0
-  - libcublas=12.4.2.65=0
-  - libcufft=11.2.0.44=0
-  - libcufile=1.10.0.4=0
-  - libcurand=10.3.6.39=0
-  - libcurl=8.8.0=hca28451_0
-  - libcusolver=11.6.0.99=0
-  - libcusparse=12.3.0.142=0
-  - libdeflate=1.20=hd590300_0
-  - libdrm=2.4.121=h4ab18f5_0
-  - libedit=3.1.20191231=he28a2e2_2
-  - libev=4.33=hd590300_2
-  - libevent=2.1.12=hf998b51_1
-  - libexpat=2.6.2=h59595ed_0
+  - ld_impl_linux-64=2.40=h41732ed_0
+  - libblas=3.9.0=16_linux64_mkl
+  - libcblas=3.9.0=16_linux64_mkl
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.9.0.20=0
+  - libcurand=10.3.5.119=0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libexpat=2.5.0=hcb278e6_1
   - libffi=3.4.2=h7f98852_5
-  - libgcc-ng=13.2.0=h77fa898_11
-  - libgettextpo=0.22.5=h59595ed_2
-  - libgettextpo-devel=0.22.5=h59595ed_2
-  - libgfortran-ng=13.2.0=h69a702a_12
-  - libgfortran5=13.2.0=h3d2ce59_12
-  - libglib=2.80.2=h8a4344b_1
-  - libgomp=13.2.0=h77fa898_11
-  - libgoogle-cloud=2.25.0=h2736e30_0
-  - libgoogle-cloud-storage=2.25.0=h3d9a0c8_0
-  - libgrpc=1.62.2=h15f2491_0
-  - libhwloc=2.10.0=default_h5622ce7_1001
-  - libiconv=1.17=hd590300_2
-  - libidn2=2.3.7=hd590300_0
-  - libjpeg-turbo=3.0.0=hd590300_1
-  - liblapack=3.9.0=22_linux64_openblas
-  - liblapacke=3.9.0=22_linux64_openblas
-  - libnghttp2=1.58.0=h47da74e_1
-  - libnpp=12.2.5.2=0
-  - libnsl=2.0.1=hd590300_0
-  - libnvfatbin=12.5.39=0
-  - libnvjitlink=12.4.99=0
-  - libnvjpeg=12.3.1.89=0
-  - libopenblas=0.3.27=pthreads_h413a1c8_0
-  - libopenvino=2024.2.0=h2da1b83_1
-  - libopenvino-auto-batch-plugin=2024.2.0=hb045406_1
-  - libopenvino-auto-plugin=2024.2.0=hb045406_1
-  - libopenvino-hetero-plugin=2024.2.0=h5c03a75_1
-  - libopenvino-intel-cpu-plugin=2024.2.0=h2da1b83_1
-  - libopenvino-intel-gpu-plugin=2024.2.0=h2da1b83_1
-  - libopenvino-intel-npu-plugin=2024.2.0=he02047a_1
-  - libopenvino-ir-frontend=2024.2.0=h5c03a75_1
-  - libopenvino-onnx-frontend=2024.2.0=h07e8aee_1
-  - libopenvino-paddle-frontend=2024.2.0=h07e8aee_1
-  - libopenvino-pytorch-frontend=2024.2.0=he02047a_1
-  - libopenvino-tensorflow-frontend=2024.2.0=h39126c6_1
-  - libopenvino-tensorflow-lite-frontend=2024.2.0=he02047a_1
-  - libopus=1.3.1=h7f98852_1
-  - libparquet=16.1.0=h6a7eafb_10_cpu
-  - libpciaccess=0.18=hd590300_0
-  - libpng=1.6.43=h2797004_0
-  - libprotobuf=4.25.3=h08a7969_0
-  - libre2-11=2023.09.01=h5a48ba9_2
-  - libsqlite=3.46.0=hde9e2c9_0
-  - libssh2=1.11.0=h0841786_0
-  - libstdcxx-ng=13.2.0=hc0a3c3a_11
-  - libtasn1=4.19.0=h166bdaf_0
-  - libthrift=0.19.0=hb90f79a_1
-  - libtiff=4.6.0=h1dd3fc0_3
-  - libtorch=2.3.1=cpu_mkl_h0bb0d08_100
-  - libunistring=0.9.10=h7f98852_0
-  - liburing=2.6=h297d8ca_0
-  - libutf8proc=2.8.0=h166bdaf_0
+  - libgcc-ng=13.2.0=h807b86a_2
+  - libgfortran-ng=13.2.0=h69a702a_5
+  - libgfortran5=13.2.0=ha4646dd_5
+  - libglib=2.78.0=hebfc3b9_0
+  - libgomp=13.2.0=h807b86a_2
+  - libiconv=1.17=h166bdaf_0
+  - liblapack=3.9.0=16_linux64_mkl
+  - libnpp=12.0.2.50=0
+  - libnsl=2.0.0=h7f98852_0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libopenblas=0.3.26=pthreads_h413a1c8_0
+  - libpng=1.6.39=h753d276_0
+  - libprotobuf=3.21.12=hfc55251_2
+  - libsqlite=3.43.0=h2797004_0
+  - libstdcxx-ng=13.2.0=h7e041cc_2
   - libuuid=2.38.1=h0b41bf4_0
-  - libuv=1.48.0=hd590300_0
-  - libva=2.22.0=hb711507_0
-  - libvpx=1.14.1=hac33072_0
-  - libwebp-base=1.4.0=hd590300_0
-  - libxcb=1.16=hd590300_0
-  - libxcrypt=4.4.36=hd590300_1
-  - libxml2=2.12.7=hc051c1a_1
-  - libzlib=1.3.1=h4ab18f5_1
-  - llvm-openmp=18.1.8=hf5423f3_0
-  - lz4-c=1.9.4=hcb278e6_0
-  - markdown=3.6=pyhd8ed1ab_0
-  - markdown-it-py=3.0.0=pyhd8ed1ab_0
-  - markupsafe=2.1.5=py310h2372a71_0
-  - matplotlib-base=3.8.4=py310hef631a5_2
-  - mccabe=0.7.0=pyhd8ed1ab_0
-  - mdurl=0.1.2=pyhd8ed1ab_0
-  - mergedeep=1.3.4=pyhd8ed1ab_0
-  - metis=5.1.1=h59595ed_2
-  - mkdocs=1.6.0=pyhd8ed1ab_0
-  - mkdocs-get-deps=0.2.0=pyhd8ed1ab_0
-  - mkdocs-material=9.5.27=pyhd8ed1ab_0
-  - mkdocs-material-extensions=1.3.1=pyhd8ed1ab_0
-  - mkl=2023.2.0=h84fe81f_50496
-  - ml_dtypes=0.3.2=py310hcc13569_0
+  - libxcb=1.15=h0b41bf4_0
+  - libxml2=2.11.5=h232c23b_1
+  - libzlib=1.2.13=hd590300_5
+  - llvm-openmp=15.0.7=h0cdce71_0
+  - markupsafe=2.1.5=py39hd1e30aa_0
+  - mkl=2022.1.0=hc2b9512_224
   - mpc=1.3.1=hfe3b2da_0
-  - mpfr=4.2.1=h9458935_1
+  - mpfr=4.2.1=h9458935_0
   - mpmath=1.3.0=pyhd8ed1ab_0
-  - multidict=6.0.5=py310h2372a71_0
-  - multiprocess=0.70.15=py310h2372a71_1
-  - munkres=1.1.4=pyh9f0ad1d_0
-  - namex=0.0.8=pyhd8ed1ab_0
-  - ncurses=6.5=h59595ed_0
-  - nettle=3.9.1=h7ab15ed_0
-  - networkx=3.3=pyhd8ed1ab_1
-  - numpy=1.26.4=py310hb13e2d6_0
-  - ocl-icd=2.3.2=hd590300_1
-  - openh264=2.4.1=h59595ed_0
-  - openjpeg=2.5.2=h488ebb8_0
-  - openssl=3.3.1=h4ab18f5_0
-  - opt_einsum=3.3.0=pyhc1e730c_2
-  - orc=2.0.1=h17fec99_1
-  - p11-kit=0.24.1=hc5aa10d_0
-  - packaging=24.1=pyhd8ed1ab_0
-  - paginate=0.5.6=pyhd8ed1ab_0
-  - pandas=2.2.2=py310hf9f9076_1
-  - pathspec=0.12.1=pyhd8ed1ab_0
-  - pcre2=10.44=h0f59acf_0
-  - pillow=10.3.0=py310hebfe307_1
-  - pip=24.0=pyhd8ed1ab_0
-  - pixman=0.43.2=h59595ed_0
-  - platformdirs=4.2.2=pyhd8ed1ab_0
-  - protobuf=4.25.3=py310ha8c1f0e_0
-  - psutil=5.9.8=py310h2372a71_0
+  - ncurses=6.4=hcb278e6_0
+  - networkx=3.2.1=pyhd8ed1ab_0
+  - ninja=1.11.1=h924138e_0
+  - openbabel=3.1.1=py39h421517d_8
+  - openssl=3.2.1=hd590300_1
+  - pcre2=10.40=hc3806b6_0
+  - pip=23.2.1=pyhd8ed1ab_0
+  - pixman=0.42.2=h59595ed_0
   - pthread-stubs=0.4=h36c2ea0_1001
-  - pugixml=1.14=h59595ed_0
-  - pyarrow=16.1.0=py310hb7f781d_3
-  - pyarrow-core=16.1.0=py310h46b3431_3_cpu
-  - pycodestyle=2.12.0=pyhd8ed1ab_0
-  - pydantic=2.7.4=pyhd8ed1ab_0
-  - pydantic-core=2.18.4=py310he421c4c_0
-  - pydantic-settings=2.3.3=pyhd8ed1ab_0
-  - pydocstyle=6.3.0=pyhd8ed1ab_0
-  - pyflakes=3.2.0=pyhd8ed1ab_0
-  - pygments=2.18.0=pyhd8ed1ab_0
-  - pymdown-extensions=10.8.1=pyhd8ed1ab_0
-  - pyparsing=2.4.7=pyhd8ed1ab_1
-  - pysocks=1.7.1=pyha2e5f31_6
-  - python=3.10.14=hd12c33a_0_cpython
-  - python-dateutil=2.9.0=pyhd8ed1ab_0
-  - python-dotenv=1.0.1=pyhd8ed1ab_0
-  - python-flatbuffers=24.3.25=pyh59ac667_0
-  - python-lmdb=1.4.1=py310hdf73078_1
-  - python-tzdata=2024.1=pyhd8ed1ab_0
-  - python-xxhash=3.4.1=py310h2372a71_0
-  - python_abi=3.10=4_cp310
-  - pytorch=2.3.1=cpu_mkl_py310h75865b9_100
-  - pytorch-cuda=12.4=hc786d27_6
-  - pytorch-mutex=1.0=cpu
-  - pytz=2024.1=pyhd8ed1ab_0
-  - pyyaml=6.0.1=py310h2372a71_1
-  - pyyaml-env-tag=0.1=pyhd8ed1ab_0
-  - re2=2023.09.01=h7f4b329_2
+  - pycparser=2.22=pyhd8ed1ab_0
+  - python=3.9.18=h0755675_0_cpython
+  - python_abi=3.9=4_cp39
+  - pytorch=2.2.2=py3.9_cuda12.1_cudnn8.9.2_0
+  - pytorch-cuda=12.1=ha16c6d3_5
+  - pytorch-mutex=1.0=cuda
+  - pyyaml=6.0.1=py39hd1e30aa_1
   - readline=8.2=h8228510_1
-  - regex=2024.5.15=py310hc51659f_0
-  - requests=2.32.3=pyhd8ed1ab_0
-  - rich=13.7.1=pyhd8ed1ab_0
-  - s2n=1.4.16=he19d79f_0
-  - scikit-learn=1.5.0=py310h981052a_1
-  - scipy=1.13.1=py310h93e2701_0
-  - setuptools=70.1.0=pyhd8ed1ab_0
-  - six=1.16.0=pyh6c4a22f_0
+  - setuptools=68.2.2=pyhd8ed1ab_0
   - sleef=3.5.1=h9b69904_2
-  - snappy=1.2.0=hdb0a2a9_1
-  - snowballstemmer=2.2.0=pyhd8ed1ab_0
-  - spglib=2.4.0=py310h02519e0_1
-  - svt-av1=2.1.0=hac33072_0
-  - sympy=1.12.1=pypyh2585a3b_103
-  - tbb=2021.12.0=h297d8ca_1
-  - tensorboard=2.16.2=pyhd8ed1ab_0
-  - tensorboard-data-server=0.7.0=py310h75e40e8_1
-  - tensorflow=2.16.1=cpu_py310h49b650b_0
-  - tensorflow-base=2.16.1=cpu_py310h224022f_0
-  - tensorflow-estimator=2.16.1=cpu_py310hc6dcfef_0
-  - termcolor=2.4.0=pyhd8ed1ab_0
-  - threadpoolctl=3.5.0=pyhc1e730c_0
-  - tk=8.6.13=noxft_h4845f30_101
-  - tomli=2.0.1=pyhd8ed1ab_0
-  - toolz=0.12.1=pyhd8ed1ab_0
-  - torchaudio=2.3.1=py310_cpu
-  - torchdata=0.7.1=py310h73bdb5a_5
-  - torchvision=0.18.1=py310_cpu
-  - tqdm=4.66.4=pyhd8ed1ab_0
-  - typing-extensions=4.12.2=hd8ed1ab_0
-  - typing_extensions=4.12.2=pyha770c72_0
-  - tzdata=2024a=h0c530f3_0
-  - unicodedata2=15.1.0=py310h2372a71_0
-  - urllib3=2.2.2=pyhd8ed1ab_0
-  - watchdog=4.0.1=py310hff52083_0
-  - wayland=1.23.0=h5291e77_0
-  - wayland-protocols=1.36=hd8ed1ab_0
-  - werkzeug=3.0.3=pyhd8ed1ab_0
+  - sympy=1.12=pypyh9d50eac_103
+  - tk=8.6.13=h2797004_0
+  - torchtriton=2.2.0=py39
+  - typing_extensions=4.10.0=pyha770c72_0
   - wheel=0.43.0=pyhd8ed1ab_1
-  - wrapt=1.16.0=py310h2372a71_0
-  - x264=1!164.3095=h166bdaf_2
-  - x265=3.5=h924138e_3
-  - xmltodict=0.13.0=pyhd8ed1ab_0
-  - xorg-fixesproto=5.0=h7f98852_1002
+  - xformers=0.0.25.post1=py39_cu12.1.0_pyt2.2.2
   - xorg-kbproto=1.0.7=h7f98852_1002
   - xorg-libice=1.1.1=hd590300_0
   - xorg-libsm=1.2.4=h7391055_0
-  - xorg-libx11=1.8.9=hb711507_1
+  - xorg-libx11=1.8.7=h8ee46fc_0
   - xorg-libxau=1.0.11=hd590300_0
   - xorg-libxdmcp=1.1.3=h7f98852_0
   - xorg-libxext=1.3.4=h0b41bf4_2
-  - xorg-libxfixes=5.0.3=h7f98852_1004
   - xorg-libxrender=0.9.11=hd590300_0
   - xorg-renderproto=0.11.1=h7f98852_1002
   - xorg-xextproto=7.3.0=h0b41bf4_1003
   - xorg-xproto=7.0.31=h7f98852_1007
-  - xxhash=0.8.2=hd590300_0
   - xz=5.2.6=h166bdaf_0
   - yaml=0.2.5=h7f98852_2
-  - yarl=1.9.4=py310h2372a71_0
-  - zipp=3.19.2=pyhd8ed1ab_0
-  - zlib=1.3.1=h4ab18f5_1
-  - zstd=1.5.6=ha6fb4c9_0
+  - zlib=1.2.13=hd590300_5
   - pip:
       - accelerate==0.31.0
+      - aiohttp==3.9.5
+      - aiosignal==1.3.1
+      - alignn==2024.4.20
+      - annotated-types==0.7.0
+      - ase==3.23.0
+      - async-timeout==4.0.3
+      - attrs==23.2.0
       - bitsandbytes==0.43.1
+      - certifi==2024.6.2
+      - chardet==3.0.4
+      - charset-normalizer==3.3.2
+      - contourpy==1.2.1
+      - cycler==0.12.1
       - datasets==2.20.0
+      - dgl==1.1.1
+      - dill==0.3.8
       - docstring-parser==0.16
-      - exceptiongroup==1.2.1
+      - eval-type-backport==0.2.0
+      - flake8==7.1.0
+      - fonttools==4.53.0
+      - frozenlist==1.4.1
       - fsspec==2024.5.0
       - huggingface-hub==0.23.4
-      - iniconfig==2.0.0
+      - idna==3.7
+      - importlib-resources==6.4.0
+      - jarvis-tools==2024.4.30
+      - joblib==1.4.2
+      - kiwisolver==1.4.5
+      - lmdb==1.4.1
+      - markdown-it-py==3.0.0
+      - matplotlib==3.9.0
+      - mccabe==0.7.0
+      - mdurl==0.1.2
+      - multidict==4.7.6
+      - multiprocess==0.70.16
+      - numpy==1.26.4
+      - packaging==24.1
+      - pandas==2.2.2
       - peft==0.11.1
-      - pluggy==1.5.0
+      - pillow==10.3.0
+      - psutil==6.0.0
+      - pyarrow==16.1.0
       - pyarrow-hotfix==0.6
-      - pytest==8.2.2
+      - pycodestyle==2.12.0
+      - pydantic==2.7.4
+      - pydantic-core==2.18.4
+      - pydantic-settings==2.3.3
+      - pydocstyle==6.3.0
+      - pyflakes==3.2.0
+      - pygments==2.18.0
+      - pyparsing==2.4.7
+      - python-dateutil==2.9.0.post0
+      - python-dotenv==1.0.1
+      - pytz==2024.1
+      - regex==2024.5.15
+      - requests==2.32.3
+      - rich==13.7.1
       - safetensors==0.4.3
+      - scikit-learn==1.5.0
+      - scipy==1.13.1
       - shtab==1.7.1
+      - six==1.16.0
+      - snowballstemmer==2.2.0
+      - spglib==2.4.0
+      - threadpoolctl==3.5.0
       - tokenizers==0.19.1
+      - toolz==0.12.1
+      - torchdata==0.7.1
+      - tqdm==4.66.4
       - transformers==4.41.2
       - trl==0.8.6
       - tyro==0.8.4
+      - tzdata==2024.1
+      - urllib3==2.2.2
+      - xmltodict==0.13.0
+      - xxhash==3.4.1
+      - yarl==1.9.4
+      - zipp==3.19.2
 prefix: atomgpt