Skip to main content

/

Docs

Nightly
v24.5
Log in

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__(out 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__(out self, spec: Optional[List[Optional[SIMD[int64, 1]]]], 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]]]]): Shape of the input, as a list of integers.
  • ​dtype (DType): Datatype of the input, from the standard library DType.

__init__(out self, spec: Optional[List[ShapeElement]], 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]]): Shape of the input, as a list of ShapeElement values.
  • ​dtype (DType): Datatype of the input, from the standard library DType.

__init__(out 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.