pykeio / ort

Fast ML inference & training for Rust with ONNX Runtime

Home Page:https://ort.pyke.io/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Enabling arena allocator

gootorov opened this issue · comments

The allocator can be switched here:

ort/src/session/mod.rs

Lines 252 to 256 in 90b8d30

/// Set the session's allocator. Defaults to [`AllocatorType::Device`].
pub fn with_allocator(mut self, allocator: AllocatorType) -> Result<Self> {
self.allocator = allocator;
Ok(self)
}

but if I'm not mistaken, self.allocator is not used anywhere, and actually building the Session just creates Allocator::default():

let allocator = Allocator::default();

Am I mistaken, or does with_allocator not do anything?

Thanks a lot for implementing it this quickly!