tch-rs icon indicating copy to clipboard operation
tch-rs copied to clipboard

Convenience methods on tch::jit::IValue to get variants

Open neilisaac opened this issue 2 years ago • 0 comments

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.

neilisaac avatar Jun 02 '22 17:06 neilisaac