Skip to main content

Mojo struct

InputSpec

Specifies a model's input shape and data type (required for TorchScript).

When loading a TorchScript model, you must specify the shape and data type for each input with an InputSpec, which you then pass to InferenceSession.load(). For example:

var batch = 1
var seqlen = 128
var input_ids_spec = TensorSpec(DType.int64, batch, seqlen)
var attention_mask_spec = TensorSpec(DType.int64, batch, seqlen)

var session = engine.InferenceSession()
var model = session.load(
"roberta.torchscript",
input_specs=List[InputSpec](
InputSpec(input_ids_spec), InputSpec(attention_mask_spec)
),
)
var batch = 1
var seqlen = 128
var input_ids_spec = TensorSpec(DType.int64, batch, seqlen)
var attention_mask_spec = TensorSpec(DType.int64, batch, seqlen)

var session = engine.InferenceSession()
var model = session.load(
"roberta.torchscript",
input_specs=List[InputSpec](
InputSpec(input_ids_spec), InputSpec(attention_mask_spec)
),
)

Implemented traits

AnyType, CollectionElement, Copyable, Movable

Methods

__init__

__init__(inout self: Self, spec: TensorSpec)

Create input specifications for one input tensor, as a TensorSpec. Only applicable for TorchScript models.

Args:

  • spec (TensorSpec): Spec for the input. This is the standard library TensorSpec.

__init__(inout self: Self, spec: Optional[List[Optional[SIMD[int64, 1]], 0]], dtype: DType)

Create specifications for one input tensor, as a list of integers. Only applicable for TorchScript models.

If an input supports dynamic shapes, use None for that dimension size.

Args:

  • spec (Optional[List[Optional[SIMD[int64, 1]], 0]]): Shape of the input, as a list of integers.
  • dtype (DType): Datatype of the input, from the standard library DType.

__init__(inout self: Self, spec: Optional[List[ShapeElement, 0]], dtype: DType)

Create specifications for one input tensor, as a list of shape elements. Only applicable for TorchScript models.

If an input supports dynamic shapes, use None or a string dimension name for that dimension size.

Args:

  • spec (Optional[List[ShapeElement, 0]]): Shape of the input, as a list of ShapeElement values.
  • dtype (DType): Datatype of the input, from the standard library DType.

__init__(inout self: Self, spec: NoneType, dtype: DType)

Create a specification for a dynamic-rank input. Only applicable for TorchScript models.

Args:

  • spec (NoneType): Always None.
  • dtype (DType): Datatype of the input, from the standard library DType.