tch-rs
tch-rs copied to clipboard
Convenience methods on tch::jit::IValue to get variants
For methods returning an IValue, we generally know what it contains, but it's awkward to consume by matching enum variants.
Simple case:
// current
let result: IValue = object.method_is(...)?;
let tensor = match result {
IValue::Tensor(t) => t,
v => { return Err(...) },
};
// desired
let tensor = object.method_is(...)?.into_tensor().map_err(...)?;
This code quickly degenerates if a model returns an IValue::Tuple of tensors for example.
Consider using EnumAsInner or hand-writing extraction methods (as_variant()
/into_variant()
) for each IValue enum variant.