Added GPU support for Windows (LaurentMazare#129)

pkubik · Dec 8, 2019 · 97bb47e · 97bb47e
1 parent ab8b59c
commit 97bb47e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/examples/basics.rs b/examples/basics.rs
@@ -1,4 +1,5 @@
 extern crate tch;
+
 use tch::{kind, Tensor};
 
 fn grad_example() {
@@ -23,4 +24,5 @@ fn main() {
     t.print();
     println!("{:?} {}", t.size(), t.double_value(&[1]));
     grad_example();
+    println!("Cuda available: {}", tch::Cuda::cudnn_is_available());
 }
diff --git a/torch-sys/build.rs b/torch-sys/build.rs
@@ -71,7 +71,7 @@ fn prepare_libtorch_dir() -> PathBuf {
 
     let device = match env::var("TORCH_CUDA_VERSION") {
         Ok(cuda_env) => match os.as_str() {
-            "linux" => cuda_env
+            "linux" | "windows" => cuda_env
                 .trim()
                 .to_lowercase()
                 .trim_start_matches("cu")
@@ -83,7 +83,7 @@ fn prepare_libtorch_dir() -> PathBuf {
                 }),
             os_str => panic!(
                 "CUDA was specified with `TORCH_CUDA_VERSION`, but pre-built \
-                 binaries with CUDA are only available for Linux, not: {}.",
+                 binaries with CUDA are only available for Linux and Windows, not: {}.",
                 os_str
             ),
         },
@@ -106,7 +106,8 @@ fn prepare_libtorch_dir() -> PathBuf {
                     TORCH_VERSION
                 ),
                 "windows" => format!(
-                    "https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-{}.zip",
+                    "https://download.pytorch.org/libtorch/{}/libtorch-win-shared-with-deps-{}.zip",
+                    device,
                     TORCH_VERSION
                 ),
                 _ => panic!("Unsupported OS"),