torch-lib-with-rust 0.1.1

No project description provided

Pytorch extension written in rust

A demo pytorch (2.3.0) package using maturin and tch.

Build in a venv

# Create and activate a venv
python -m venv .venv
source .venv/bin/activate
# Install maturin
pip install maturin
# Build the torch_lib_with_rust package with maturin
LIBTORCH_USE_PYTORCH=1 maturin develop

Run

import torch
import torch_lib_with_rust

a=torch.randn(16).to("cuda")

print(a)
print(torch_lib_with_rust.add_one(a))