Update tch instructions (tracel-ai#2844)

* Update tch instructions

* Add windows note

Author: laggui

crates/burn-tch/README.md

@@ -17,7 +17,7 @@ The backend supports CPU (multithreaded), [CUDA](https://pytorch.org/docs/stable
 [`tch-rs`](https://github.com/LaurentMazare/tch-rs) requires the C++ PyTorch library (LibTorch) to
 be available on your system.
 
-By default, the CPU distribution is installed for LibTorch v2.2.0 as required by `tch-rs`.
+By default, the CPU distribution is installed for LibTorch v2.6.0 as required by `tch-rs`.
 
 <details>
 <summary><strong>CUDA</strong></summary>
@@ -26,20 +26,25 @@ To install the latest compatible CUDA distribution, set the `TORCH_CUDA_VERSION`
 variable before the `tch-rs` dependency is retrieved with `cargo`.
 
 ```shell
-export TORCH_CUDA_VERSION=cu121
+export TORCH_CUDA_VERSION=cu124
 ```
 
 On Windows:
 
 ```powershell
-$Env:TORCH_CUDA_VERSION = "cu121"
+$Env:TORCH_CUDA_VERSION = "cu124"
 ```
 
+> Note: `tch` doesn't expose the downloaded libtorch directory on Windows when using the automatic
+> download feature, so the `torch_cuda.dll` cannot be detected properly during build. In this case,
+> you can set the `LIBTORCH` environment variable to point to the `libtorch/` folder in `torch-sys`
+> `OUT_DIR` (or move the downloaded lib to a different folder and point to it).
+
 For example, running the validation sample for the first time could be done with the following
 commands:
 
 ```shell
-export TORCH_CUDA_VERSION=cu121
+export TORCH_CUDA_VERSION=cu124
 cargo run --bin cuda --release
 ```
 
@@ -88,7 +93,7 @@ platform.
 First, download the LibTorch CPU distribution.
 
 ```shell
-wget -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcpu.zip
+wget -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcpu.zip
 unzip libtorch.zip
 ```
 
@@ -108,7 +113,7 @@ export LD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH
 First, download the LibTorch CPU distribution.
 
 ```shell
-wget -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-macos-x86_64-2.2.0.zip
+wget -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-macos-x86_64-2.6.0.zip
 unzip libtorch.zip
 ```
 
@@ -128,7 +133,7 @@ export DYLD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$DYLD_LIBRARY_PATH
 First, download the LibTorch CPU distribution.
 
 ```powershell
-wget https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.2.0%2Bcpu.zip -OutFile libtorch.zip
+wget https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.6.0%2Bcpu.zip -OutFile libtorch.zip
 Expand-Archive libtorch.zip
 ```
 
@@ -144,62 +149,17 @@ $Env:Path += ";/absolute/path/to/libtorch/"
 
 #### CUDA
 
-LibTorch 2.2.0 currently includes binary distributions with CUDA 11.8 or 12.1 runtimes. The manual
-installation instructions are detailed below.
-
-**CUDA 11.8**
-
-<details open>
-<summary><strong>🐧 Linux</strong></summary>
-
-First, download the LibTorch CUDA 11.8 distribution.
-
-```shell
-wget -O libtorch.zip https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcu118.zip
-unzip libtorch.zip
-```
-
-Then, point to that installation using the `LIBTORCH` and `LD_LIBRARY_PATH` environment variables
-before building `burn-tch` or a crate which depends on it.
-
-```shell
-export LIBTORCH=/absolute/path/to/libtorch/
-export LD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH
-```
-
-**Note:** make sure your CUDA installation is in your `PATH` and `LD_LIBRARY_PATH`.
-
-</details><br>
-
-<details>
-<summary><strong>🪟 Windows</strong></summary>
-
-First, download the LibTorch CUDA 11.8 distribution.
-
-```powershell
-wget https://download.pytorch.org/libtorch/cu118/libtorch-win-shared-with-deps-2.2.0%2Bcu118.zip -OutFile libtorch.zip
-Expand-Archive libtorch.zip
-```
-
-Then, set the `LIBTORCH` environment variable and append the library to your path as with the
-PowerShell commands below before building `burn-tch` or a crate which depends on it.
-
-```powershell
-$Env:LIBTORCH = "/absolute/path/to/libtorch/"
-$Env:Path += ";/absolute/path/to/libtorch/"
-```
-
-</details><br>
-
-**CUDA 12.1**
+LibTorch 2.6.0 currently includes binary distributions with CUDA 11.8, 12.4 or 12.6 runtimes. The
+manual installation instructions are detailed below for CUDA 12.6, but can be applied to the other
+CUDA versions by replacing `cu126` with the corresponding version string (e.g., `cu118` or `cu124`).
 
 <details open>
 <summary><strong>🐧 Linux</strong></summary>
 
-First, download the LibTorch CUDA 12.1 distribution.
+First, download the LibTorch CUDA 12.6 distribution.
 
 ```shell
-wget -O libtorch.zip https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcu121.zip
+wget -O libtorch.zip https://download.pytorch.org/libtorch/cu126/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcu126.zip
 unzip libtorch.zip
 ```
 
@@ -218,10 +178,10 @@ export LD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH
 <details>
 <summary><strong>🪟 Windows</strong></summary>
 
-First, download the LibTorch CUDA 12.1 distribution.
+First, download the LibTorch CUDA 12.6 distribution.
 
 ```powershell
-wget https://download.pytorch.org/libtorch/cu121/libtorch-win-shared-with-deps-2.2.0%2Bcu121.zip -OutFile libtorch.zip
+wget https://download.pytorch.org/libtorch/cu126/libtorch-win-shared-with-deps-2.6.0%2Bcu126.zip -OutFile libtorch.zip
 Expand-Archive libtorch.zip
 ```
 
@@ -243,13 +203,13 @@ is to use a PyTorch installation. This requires a Python installation.
 _Note: MPS acceleration is available on MacOS 12.3+._
 
 ```shell
-pip install torch==2.2.0 numpy==1.26.4 setuptools
+pip install torch==2.6.0 numpy==1.26.4 setuptools
 export LIBTORCH_USE_PYTORCH=1
 export DYLD_LIBRARY_PATH=/path/to/pytorch/lib:$DYLD_LIBRARY_PATH
 ```
 
-**Note:** if `venv` is used, it should be activated during coding and building,
-or the compiler may not work properly.
+**Note:** if `venv` is used, it should be activated during coding and building, or the compiler may
+not work properly.
 
 ## Example Usage
 
@@ -263,7 +223,8 @@ For a more complete example using the `tch` backend, take a loot at the
 
 Try `.cargo/config.toml` ([cargo book](https://doc.rust-lang.org/cargo/reference/config.html#env)).
 
-Instead of setting the environments in your shell, you can manually add them to your `.cargo/config.toml`:
+Instead of setting the environments in your shell, you can manually add them to your
+`.cargo/config.toml`:
 
 ```toml
 [env]
@@ -281,4 +242,5 @@ LD_LIBRARY_PATH = "/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH"
 LIBTORCH = "/absolute/path/to/libtorch/libtorch"
 EOF
 ```
+
 This will automatically include the old `LD_LIBRARY_PATH` value in the new one.

examples/custom-image-dataset/README.md

@@ -54,7 +54,7 @@ achieves 70-80% accuracy on the test set after just 30 epochs.
 Run it with the Torch GPU backend:
 
 ```sh
-export TORCH_CUDA_VERSION=cu121
+export TORCH_CUDA_VERSION=cu124
 cargo run --example custom-image-dataset --release --features tch-gpu
 ```
 

examples/mnist/README.md

@@ -17,7 +17,7 @@ cargo run --example mnist --release --features ndarray                # CPU NdAr
 cargo run --example mnist --release --features ndarray-blas-openblas  # CPU NdArray Backend - f32 - blas with openblas
 cargo run --example mnist --release --features ndarray-blas-netlib    # CPU NdArray Backend - f32 - blas with netlib
 echo "Using tch backend"
-export TORCH_CUDA_VERSION=cu121                                       # Set the cuda version
+export TORCH_CUDA_VERSION=cu124                                       # Set the cuda version
 cargo run --example mnist --release --features tch-gpu                # GPU Tch Backend - f32
 cargo run --example mnist --release --features tch-cpu                # CPU Tch Backend - f32
 echo "Using wgpu backend"

examples/modern-lstm/README.md

@@ -27,7 +27,7 @@ cargo run --example lstm-train --release --features cuda
 cargo run --example lstm-train --release --features wgpu
 
 # Tch GPU backend
-export TORCH_CUDA_VERSION=cu121 # Set the cuda version
+export TORCH_CUDA_VERSION=cu124 # Set the cuda version
 cargo run --example lstm-train --release --features tch-gpu
 
 # Tch CPU backend

examples/simple-regression/README.md

@@ -26,7 +26,7 @@ cargo run --example regression --release --features ndarray                # CPU
 cargo run --example regression --release --features ndarray-blas-openblas  # CPU NdArray Backend - f32 - blas with openblas
 cargo run --example regression --release --features ndarray-blas-netlib    # CPU NdArray Backend - f32 - blas with netlib
 echo "Using tch backend"
-export TORCH_CUDA_VERSION=cu121                                            # Set the cuda version
+export TORCH_CUDA_VERSION=cu124                                            # Set the cuda version
 cargo run --example regression --release --features tch-gpu                # GPU Tch Backend - f32
 cargo run --example regression --release --features tch-cpu                # CPU Tch Backend - f32
 echo "Using wgpu backend"

examples/text-classification/README.md

@@ -29,7 +29,7 @@ cd burn
 # Use the --release flag to really speed up training.
 # Use the f16 feature if your CUDA device supports FP16 (half precision) operations. May not work well on every device.
 
-export TORCH_CUDA_VERSION=cu121  # Set the cuda version (CUDA users)
+export TORCH_CUDA_VERSION=cu124  # Set the cuda version (CUDA users)
 
 # AG News
 cargo run --example ag-news-train --release --features tch-gpu   # Train on the ag news dataset

examples/text-generation/README.md

@@ -14,7 +14,7 @@ git clone https://github.com/tracel-ai/burn.git
 cd burn
 
 # Use the --release flag to really speed up training.
-export TORCH_CUDA_VERSION=cu121
+export TORCH_CUDA_VERSION=cu124
 cargo run --example text-generation --release
 ```
 

examples/wgan/README.md

@@ -18,7 +18,7 @@ cargo run --example wgan-mnist --release --features cuda
 cargo run --example wgan-mnist --release --features wgpu
 
 # Tch GPU backend
-export TORCH_CUDA_VERSION=cu121 # Set the cuda version
+export TORCH_CUDA_VERSION=cu124 # Set the cuda version
 cargo run --example wgan-mnist --release --features tch-gpu
 
 # Tch CPU backend