diff --git a/env.yaml b/env.yaml
new file mode 100644
index 0000000..3c1c917
--- /dev/null
+++ b/env.yaml
@@ -0,0 +1,310 @@
+name: 626_pytorch_lira
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_gnu
+  - alsa-lib=1.2.3.2=h166bdaf_0
+  - anyio=4.5.0=pyhd8ed1ab_0
+  - appdirs=1.4.4=pyh9f0ad1d_0
+  - argon2-cffi=23.1.0=pyhd8ed1ab_0
+  - argon2-cffi-bindings=21.2.0=py38h01eb140_4
+  - arrow=1.3.0=pyhd8ed1ab_0
+  - asttokens=2.4.1=pyhd8ed1ab_0
+  - async-lru=2.0.4=pyhd8ed1ab_0
+  - attrs=24.2.0=pyh71513ae_0
+  - babel=2.16.0=pyhd8ed1ab_0
+  - backcall=0.2.0=pyh9f0ad1d_0
+  - beautifulsoup4=4.12.3=pyha770c72_0
+  - blas=1.0=mkl
+  - bleach=6.1.0=pyhd8ed1ab_0
+  - brotli=1.1.0=hd590300_1
+  - brotli-bin=1.1.0=hd590300_1
+  - brotli-python=1.1.0=py38h17151c0_1
+  - bzip2=1.0.8=h4bc722e_7
+  - ca-certificates=2024.8.30=hbcca054_0
+  - cached-property=1.5.2=hd8ed1ab_1
+  - cached_property=1.5.2=pyha770c72_1
+  - certifi=2024.8.30=pyhd8ed1ab_0
+  - cffi=1.14.6=py38ha65f79e_0
+  - charset-normalizer=3.4.0=pyhd8ed1ab_0
+  - click=8.1.7=unix_pyh707e725_0
+  - colorama=0.4.6=pyhd8ed1ab_0
+  - comm=0.2.2=pyhd8ed1ab_0
+  - contourpy=1.1.1=py38h7f3f72f_1
+  - 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.6.77=0
+  - cuda-runtime=12.1.0=0
+  - cuda-version=12.6=3
+  - cycler=0.12.1=pyhd8ed1ab_0
+  - dataclasses=0.8=pyhc8e2a94_3
+  - dbus=1.13.6=h48d8840_2
+  - debugpy=1.8.5=py38h6d02427_0
+  - decorator=5.1.1=pyhd8ed1ab_0
+  - defusedxml=0.7.1=pyhd8ed1ab_0
+  - docker-pycreds=0.4.0=py_0
+  - entrypoints=0.4=pyhd8ed1ab_0
+  - exceptiongroup=1.2.2=pyhd8ed1ab_0
+  - executing=2.1.0=pyhd8ed1ab_0
+  - expat=2.6.4=h5888daf_0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.16.1=pyhd8ed1ab_0
+  - fontconfig=2.14.2=h14ed4e7_0
+  - fonttools=4.53.1=py38h2019614_0
+  - fqdn=1.5.1=pyhd8ed1ab_0
+  - freetype=2.12.1=h267a509_2
+  - fsspec=2024.10.0=pyhff2d567_0
+  - functorch=2.0.0=pyhd8ed1ab_0
+  - gettext=0.22.5=he02047a_3
+  - gettext-tools=0.22.5=he02047a_3
+  - giflib=5.2.2=hd590300_0
+  - gitdb=4.0.11=pyhd8ed1ab_0
+  - gitpython=3.1.43=pyhd8ed1ab_0
+  - glib=2.68.4=h9c3ff4c_0
+  - glib-tools=2.68.4=h9c3ff4c_0
+  - gmp=6.3.0=hac33072_2
+  - gmpy2=2.1.5=py38h6a1700d_1
+  - gnutls=3.6.13=h85f3911_1
+  - gst-plugins-base=1.18.5=hf529b03_0
+  - gstreamer=1.18.5=h76c114f_0
+  - h11=0.14.0=pyhd8ed1ab_0
+  - h2=4.1.0=pyhd8ed1ab_0
+  - hpack=4.0.0=pyh9f0ad1d_0
+  - httpcore=1.0.7=pyh29332c3_1
+  - httpx=0.27.2=pyhd8ed1ab_0
+  - hyperframe=6.0.1=pyhd8ed1ab_0
+  - icu=68.2=h9c3ff4c_0
+  - idna=3.10=pyhd8ed1ab_0
+  - importlib-metadata=8.5.0=pyha770c72_0
+  - importlib-resources=6.4.5=pyhd8ed1ab_0
+  - importlib_resources=6.4.5=pyhd8ed1ab_0
+  - iniconfig=2.0.0=pyhd8ed1ab_0
+  - intel-openmp=2022.1.0=h9e868ea_3769
+  - ipykernel=6.29.5=pyh3099207_0
+  - ipython=8.12.2=pyh41d4057_0
+  - ipywidgets=8.1.5=pyhd8ed1ab_0
+  - isoduration=20.11.0=pyhd8ed1ab_0
+  - jedi=0.19.1=pyhd8ed1ab_0
+  - jinja2=3.1.4=pyhd8ed1ab_0
+  - joblib=1.4.2=pyhd8ed1ab_0
+  - jpeg=9e=h166bdaf_2
+  - json5=0.9.25=pyhd8ed1ab_0
+  - jsonpointer=3.0.0=py38h578d9bd_0
+  - jsonschema=4.23.0=pyhd8ed1ab_0
+  - jsonschema-specifications=2024.10.1=pyhd8ed1ab_0
+  - jsonschema-with-format-nongpl=4.23.0=hd8ed1ab_0
+  - jupyter=1.1.1=pyhd8ed1ab_0
+  - jupyter-lsp=2.2.5=pyhd8ed1ab_0
+  - jupyter_client=8.6.3=pyhd8ed1ab_0
+  - jupyter_console=6.6.3=pyhd8ed1ab_0
+  - jupyter_core=5.7.2=pyh31011fe_1
+  - jupyter_events=0.10.0=pyhd8ed1ab_0
+  - jupyter_server=2.14.2=pyhd8ed1ab_0
+  - jupyter_server_terminals=0.5.3=pyhd8ed1ab_0
+  - jupyterlab=4.3.0=pyhd8ed1ab_0
+  - jupyterlab_pygments=0.3.0=pyhd8ed1ab_1
+  - jupyterlab_server=2.27.3=pyhd8ed1ab_0
+  - jupyterlab_widgets=3.0.13=pyhd8ed1ab_0
+  - keyutils=1.6.1=h166bdaf_0
+  - kiwisolver=1.4.5=py38h7f3f72f_1
+  - krb5=1.19.3=h3790be6_0
+  - lame=3.100=h166bdaf_1003
+  - lcms2=2.15=hfd0df8a_0
+  - ld_impl_linux-64=2.43=h712a8e2_2
+  - lerc=4.0.0=h27087fc_0
+  - libabseil=20240116.2=cxx17_he02047a_1
+  - libasprintf=0.22.5=he8f35ee_3
+  - libasprintf-devel=0.22.5=he8f35ee_3
+  - libblas=3.9.0=16_linux64_mkl
+  - libbrotlicommon=1.1.0=hd590300_1
+  - libbrotlidec=1.1.0=hd590300_1
+  - libbrotlienc=1.1.0=hd590300_1
+  - libcblas=3.9.0=16_linux64_mkl
+  - libclang=11.1.0=default_ha53f305_1
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.11.1.6=0
+  - libcurand=10.3.7.77=0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libdeflate=1.17=h0b41bf4_0
+  - libedit=3.1.20191231=he28a2e2_2
+  - libevent=2.1.10=h9b69904_4
+  - libexpat=2.6.4=h5888daf_0
+  - libffi=3.3=h58526e2_2
+  - libgcc=14.2.0=h77fa898_1
+  - libgcc-ng=14.2.0=h69a702a_1
+  - libgettextpo=0.22.5=he02047a_3
+  - libgettextpo-devel=0.22.5=he02047a_3
+  - libgfortran=14.2.0=h69a702a_1
+  - libgfortran-ng=14.2.0=h69a702a_1
+  - libgfortran5=14.2.0=hd5240d6_1
+  - libglib=2.68.4=h3e27bee_0
+  - libgomp=14.2.0=h77fa898_1
+  - libiconv=1.17=hd590300_2
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - liblapack=3.9.0=16_linux64_mkl
+  - libllvm11=11.1.0=he0ac6c6_5
+  - libnpp=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libogg=1.3.5=h4ab18f5_0
+  - libopus=1.3.1=h7f98852_1
+  - libpng=1.6.43=h2797004_0
+  - libpq=13.8=hd77ab85_0
+  - libprotobuf=4.25.3=h08a7969_0
+  - libsodium=1.0.18=h36c2ea0_1
+  - libsqlite=3.46.0=hde9e2c9_0
+  - libstdcxx=14.2.0=hc0a3c3a_1
+  - libstdcxx-ng=14.2.0=h4852527_1
+  - libtiff=4.5.0=h6adf6a1_2
+  - libuuid=2.38.1=h0b41bf4_0
+  - libvorbis=1.3.7=h9c3ff4c_0
+  - libwebp=1.3.2=h11a3e52_0
+  - libwebp-base=1.3.2=hd590300_1
+  - libxcb=1.13=h7f98852_1004
+  - libxkbcommon=1.0.3=he3ba5ed_0
+  - libxml2=2.9.12=h72842e0_0
+  - libzlib=1.2.13=h4ab18f5_6
+  - lightning-bolts=0.6.0.post1=pyhd8ed1ab_0
+  - lightning-utilities=0.11.8=pyhd8ed1ab_0
+  - llvm-openmp=15.0.7=h0cdce71_0
+  - markupsafe=2.1.5=py38h01eb140_0
+  - matplotlib=3.7.3=py38h578d9bd_0
+  - matplotlib-base=3.7.3=py38h58ed7fa_0
+  - matplotlib-inline=0.1.7=pyhd8ed1ab_0
+  - mistune=3.0.2=pyhd8ed1ab_0
+  - mkl=2022.1.0=hc2b9512_224
+  - mpc=1.3.1=h24ddda3_1
+  - mpfr=4.2.1=h90cbb55_3
+  - mpmath=1.3.0=pyhd8ed1ab_0
+  - munkres=1.1.4=pyh9f0ad1d_0
+  - mysql-common=8.0.32=h14678bc_0
+  - mysql-libs=8.0.32=h54cf53e_0
+  - nbclient=0.10.1=pyhd8ed1ab_0
+  - nbconvert-core=7.16.4=pyhd8ed1ab_1
+  - nbformat=5.10.4=pyhd8ed1ab_0
+  - ncurses=6.5=he02047a_1
+  - nest-asyncio=1.6.0=pyhd8ed1ab_0
+  - nettle=3.6=he412f7d_0
+  - networkx=3.1=pyhd8ed1ab_0
+  - notebook=7.0.6=py38h06a4308_0
+  - notebook-shim=0.2.4=pyhd8ed1ab_0
+  - nspr=4.36=h5888daf_0
+  - nss=3.100=hca3bf56_0
+  - numpy=1.24.4=py38h59b608b_0
+  - opacus=1.5.2=pyhd8ed1ab_0
+  - openh264=2.1.1=h780b84a_0
+  - openjpeg=2.5.0=hfec8fc6_2
+  - openssl=1.1.1w=hd590300_0
+  - opt_einsum=3.4.0=pyhd8ed1ab_0
+  - overrides=7.7.0=pyhd8ed1ab_0
+  - packaging=24.2=pyhff2d567_1
+  - pandas=2.0.3=py38h01efb38_1
+  - pandocfilters=1.5.0=pyhd8ed1ab_0
+  - parso=0.8.4=pyhd8ed1ab_0
+  - pcre=8.45=h9c3ff4c_0
+  - pexpect=4.9.0=pyhd8ed1ab_0
+  - pickleshare=0.7.5=py_1003
+  - pillow=9.4.0=py38hde6dc18_1
+  - pip=24.3.1=pyh8b19718_0
+  - pkgutil-resolve-name=1.3.10=pyhd8ed1ab_1
+  - platformdirs=4.3.6=pyhd8ed1ab_0
+  - pluggy=1.5.0=pyhd8ed1ab_0
+  - pooch=1.8.2=pyhd8ed1ab_0
+  - prometheus_client=0.21.0=pyhd8ed1ab_0
+  - prompt-toolkit=3.0.48=pyha770c72_0
+  - prompt_toolkit=3.0.48=hd8ed1ab_0
+  - protobuf=4.25.3=py38hb5c7596_0
+  - psutil=6.0.0=py38hfb59056_0
+  - pthread-stubs=0.4=hb9d3cd8_1002
+  - ptyprocess=0.7.0=pyhd3deb0d_0
+  - pure_eval=0.2.3=pyhd8ed1ab_0
+  - pycparser=2.22=pyhd8ed1ab_0
+  - pygments=2.18.0=pyhd8ed1ab_0
+  - pyparsing=3.1.4=pyhd8ed1ab_0
+  - pyqt=5.12.3=py38h578d9bd_8
+  - pyqt-impl=5.12.3=py38h0ffb2e6_8
+  - pyqt5-sip=4.19.18=py38h709712a_8
+  - pyqtchart=5.12=py38h7400c14_8
+  - pyqtwebengine=5.12.1=py38h7400c14_8
+  - pysocks=1.7.1=pyha2e5f31_6
+  - pytest=8.3.3=pyhd8ed1ab_0
+  - python=3.8.6=hffdb5ce_5_cpython
+  - python-dateutil=2.9.0=pyhd8ed1ab_0
+  - python-fastjsonschema=2.20.0=pyhd8ed1ab_0
+  - python-json-logger=2.0.7=pyhd8ed1ab_0
+  - python-tzdata=2024.2=pyhd8ed1ab_0
+  - python_abi=3.8=5_cp38
+  - pytorch=2.4.1=py3.8_cuda12.1_cudnn9.1.0_0
+  - pytorch-cuda=12.1=ha16c6d3_6
+  - pytorch-lightning=2.4.0=pyhd8ed1ab_0
+  - pytorch-mutex=1.0=cuda
+  - pytz=2024.2=pyhd8ed1ab_0
+  - pyyaml=6.0.2=py38h2019614_0
+  - pyzmq=26.2.0=py38h6c80b9a_0
+  - qt=5.12.9=hda022c4_4
+  - readline=8.2=h8228510_1
+  - referencing=0.35.1=pyhd8ed1ab_0
+  - requests=2.32.3=pyhd8ed1ab_0
+  - rfc3339-validator=0.1.4=pyhd8ed1ab_0
+  - rfc3986-validator=0.1.1=pyh9f0ad1d_0
+  - rpds-py=0.20.0=py38h4005ec7_0
+  - scikit-learn=1.3.2=py38ha25d942_2
+  - scipy=1.10.1=py38h59b608b_3
+  - send2trash=1.8.3=pyh0d859eb_0
+  - sentry-sdk=2.19.0=pyhd8ed1ab_0
+  - setproctitle=1.3.3=py38h01eb140_0
+  - setuptools=75.3.0=pyhd8ed1ab_0
+  - six=1.16.0=pyh6c4a22f_0
+  - smmap=5.0.0=pyhd8ed1ab_0
+  - sniffio=1.3.1=pyhd8ed1ab_0
+  - soupsieve=2.5=pyhd8ed1ab_1
+  - sqlite=3.46.0=h6d4b2fc_0
+  - stack_data=0.6.2=pyhd8ed1ab_0
+  - sympy=1.13.3=pypyh2585a3b_103
+  - terminado=0.18.1=pyh0d859eb_0
+  - threadpoolctl=3.5.0=pyhc1e730c_0
+  - tinycss2=1.4.0=pyhd8ed1ab_0
+  - tk=8.6.13=noxft_h4845f30_101
+  - tomli=2.0.2=pyhd8ed1ab_0
+  - torchaudio=2.4.1=py38_cu121
+  - torchmetrics=1.5.2=pyhe5570ce_0
+  - torchtriton=3.0.0=py38
+  - torchvision=0.20.0=py38_cu121
+  - tornado=6.4.1=py38hfb59056_0
+  - tqdm=4.67.1=pyhd8ed1ab_0
+  - traitlets=5.14.3=pyhd8ed1ab_0
+  - types-python-dateutil=2.9.0.20241003=pyhff2d567_0
+  - typing-extensions=4.12.2=hd8ed1ab_0
+  - typing_extensions=4.12.2=pyha770c72_0
+  - typing_utils=0.1.0=pyhd8ed1ab_0
+  - unicodedata2=15.1.0=py38h01eb140_0
+  - uri-template=1.3.0=pyhd8ed1ab_0
+  - urllib3=2.2.3=pyhd8ed1ab_0
+  - wandb=0.16.6=pyhd8ed1ab_1
+  - wcwidth=0.2.13=pyhd8ed1ab_0
+  - webcolors=24.8.0=pyhd8ed1ab_0
+  - webencodings=0.5.1=pyhd8ed1ab_2
+  - websocket-client=1.8.0=pyhd8ed1ab_0
+  - wheel=0.45.1=pyhd8ed1ab_0
+  - widgetsnbextension=4.0.13=pyhd8ed1ab_0
+  - xorg-libxau=1.0.11=hb9d3cd8_1
+  - xorg-libxdmcp=1.1.5=hb9d3cd8_0
+  - xz=5.2.6=h166bdaf_0
+  - yaml=0.2.5=h7f98852_2
+  - zeromq=4.3.5=h59595ed_1
+  - zipp=3.21.0=pyhd8ed1ab_0
+  - zlib=1.2.13=h4ab18f5_6
+  - zstandard=0.23.0=py38h62bed22_0
+  - zstd=1.5.6=ha6fb4c9_0
+  - pip:
+      - pyvacy==0.0.32
+      - torchsummary==1.5.1