Burn Torch backend
This crate provides a Torch backend for Burn utilizing the tch-rs crate, which offers a Rust interface to the PyTorch C++ API.
The backend supports CPU (multithreaded), CUDA (multiple GPUs), and MPS devices (MacOS).
#[cfg(feature = "tch-gpu")]
mod tch_gpu {
use burn_autodiff::Autodiff;
use burn_tch::{LibTorch, LibTorchDevice};
use mnist::training;
pub fn run() {
#[cfg(not(target_os = "macos"))]
let device = LibTorchDevice::Cuda(0);
#[cfg(target_os = "macos")]
let device = LibTorchDevice::Mps;
training::run::<Autodiff<LibTorch<f32>>>(device);
}
}
#[cfg(feature = "tch-cpu")]
mod tch_cpu {
use burn_autodiff::Autodiff;
use burn_tch::{LibTorch, LibTorchDevice};
use mnist::training;
pub fn run() {
let device = LibTorchDevice::Cpu;
training::run::<Autodiff<LibTorch<f32>>>(device);
}
}