From b63fc70e9b671988f67aec1fbc62bbd01f830b25 Mon Sep 17 00:00:00 2001
From: Raul <raulppelaez@gmail.com>
Date: Tue, 18 Apr 2023 11:16:44 +0200
Subject: [PATCH] Support for pytorch 2.0 (#94)

* Update environments to support torch 2.0

* Make latest CUDA version in CI be 11.7, since the CUDA workflow provider does not
include 11.8

* Update ci

* Update ci

* Update ci

* Update ci

* Add latest torchani, compatible with pytorch 2

* update ci

* update ci

* Update ci

* Address raimis comments

* Fix cuda

* Remove = in environment.yml
---
 .github/workflows/ci.yml | 22 ++++++++++++++--------
 environment.yml          | 10 +++++-----
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index aa476ff0..760f69bd 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -29,6 +29,7 @@ jobs:
             gcc: "8.5.*"
             nvcc: "10.2"
             python: "3.8.*"
+            torchani: "2.2.*"
             pytorch: "1.11.*"
 
           # Older supported versions
@@ -38,33 +39,36 @@ jobs:
             gcc: "10.3.*"
             nvcc: "11.2"
             python: "3.9.*"
+            torchani: "2.2.*"
             pytorch: "1.12.*"
 
           # Latest supported versions (with CUDA)
-          - name: Linux (CUDA 11.7, Python 3.10, PyTorch 1.13)
+          - name: Linux (CUDA 11.8, Python 3.10, PyTorch 2.0)
             enable_cuda: true
-            cuda: "11.7.0"
+            cuda: "11.8.0"
             gcc: "10.3.*"
-            nvcc: "11.7"
+            nvcc: "11.8"
             python: "3.10.*"
-            pytorch: "1.13.*"
+            torchani: "2.2.*"
+            pytorch: "2.0.*"
 
           # Latest supported versions (without CUDA)
-          - name: Linux (no CUDA, Python 3.10, PyTorch 1.13)
+          - name: Linux (no CUDA, Python 3.10, PyTorch 2.0)
             enable_cuda: false
             gcc: "10.3.*"
             python: "3.10.*"
-            pytorch: "1.13.*"
+            pytorch: "2.0.*"
+            torchani: "2.2.*"
 
     steps:
     - name: Check out
       uses: actions/checkout@v2
 
     - name: Install CUDA Toolkit
-      uses: Jimver/cuda-toolkit@v0.2.8
+      uses: Jimver/cuda-toolkit@v0.2.10
       with:
         cuda: ${{ matrix.cuda }}
-        linux-local-args: '["--toolkit", "--override"]' # Need to install CUDA 10.2
+        linux-local-args: '["--toolkit", "--override"]' 
       if: ${{ matrix.enable_cuda }}
 
     - name: Install Miniconda
@@ -79,6 +83,7 @@ jobs:
       run: |
         sed -i -e "/cudatoolkit/c\  - cudatoolkit ${{ matrix.cuda }}" \
                -e "/gxx_linux-64/c\  - gxx_linux-64 ${{ matrix.gcc }}" \
+               -e "/torchani/c\  - torchani ${{ matrix.torchani }}" \
                -e "/nvcc_linux-64/c\  - nvcc_linux-64 ${{ matrix.nvcc }}" \
                -e "/python/c\  - python ${{ matrix.python }}" \
                -e "/pytorch-gpu/c\  - pytorch-gpu ${{ matrix.pytorch }}" \
@@ -89,6 +94,7 @@ jobs:
       run: |
         sed -i -e "/cudatoolkit/c\  # - cudatoolkit" \
                -e "/gxx_linux-64/c\  - gxx_linux-64 ${{ matrix.gcc }}" \
+               -e "/torchani/c\  - torchani ${{ matrix.torchani }}" \
                -e "/nvcc_linux-64/c\  # - nvcc_linux-64" \
                -e "/python/c\  - python ${{ matrix.python }}" \
                -e "/pytorch-gpu/c\  - pytorch-cpu ${{ matrix.pytorch }}" \
diff --git a/environment.yml b/environment.yml
index e2caabf9..9db9ae23 100644
--- a/environment.yml
+++ b/environment.yml
@@ -2,13 +2,13 @@ channels:
   - conda-forge
 dependencies:
   - cmake >=3.20
-  - cudatoolkit 11.2.2
+  - cudatoolkit 11.8.*
   - gxx_linux-64 10.3.*
   - make
   - mdtraj
-  - nvcc_linux-64 11.2
-  - torchani 2.2.2
+  - nvcc_linux-64 11.8
+  - torchani 2.2.*
   - pytest
   - python 3.10.*
-  - pytorch-gpu 1.12.*
-  - sysroot_linux-64 2.17
\ No newline at end of file
+  - pytorch-gpu 2.0.*
+  - sysroot_linux-64 2.17