tch 0.1.6

Rust wrappers for the PyTorch C++ api (libtorch).
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
extern crate tch;
use tch::Tensor;

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let device = match a.iter().map(|x| x.as_str()).collect::<Vec<_>>().as_slice() {
        [_] => tch::Device::Cpu,
        [_, "cpu"] => tch::Device::Cpu,
        [_, "gpu"] => tch::Device::Cuda(0),
        _ => panic!("usage: main cpu|gpu"),
    };
    let slice = vec![0; 1_000_000];
    for i in 1..1_000_000 {
        let t = Tensor::of_slice(&slice).to_device(device);
        println!("{} {:?}", i, t.size())
    }
}