tch-save-image.rs

use failure::Fallible;
use ndarray::Array;
use std::convert::TryFrom;
use tch::Tensor;

fn main() -> Fallible<()> {
    let tensor = Tensor::try_from(Array::from_shape_fn(
        (3, 300, 200),
        |(channel, _row, _col)| match channel {
            0 => 1f32,
            1 | 2 => 0.0,
            _ => unreachable!(),
        },
    ))?;
    tch::vision::image::save(&tensor, "test.jpg")?;
    Ok(())
}