File: _experimental.py

package info (click to toggle)
pytorch-cuda 2.6.0%2Bdfsg-7
  • links: PTS, VCS
  • area: contrib
  • in suites: forky, sid, trixie
  • size: 161,620 kB
  • sloc: python: 1,278,832; cpp: 900,322; ansic: 82,710; asm: 7,754; java: 3,363; sh: 2,811; javascript: 2,443; makefile: 597; ruby: 195; xml: 84; objc: 68
file content (27 lines) | stat: -rw-r--r-- 1,017 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
"""Experimental classes and functions used by ONNX export."""

import dataclasses
from typing import Mapping, Optional, Sequence, Set, Type, Union

import torch
import torch._C._onnx as _C_onnx


@dataclasses.dataclass
class ExportOptions:
    """Arguments used by :func:`torch.onnx.export`."""

    # TODO(justinchuby): Deprecate and remove this class.

    export_params: bool = True
    verbose: bool = False
    training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL
    input_names: Optional[Sequence[str]] = None
    output_names: Optional[Sequence[str]] = None
    operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX
    opset_version: Optional[int] = None
    do_constant_folding: bool = True
    dynamic_axes: Optional[Mapping[str, Union[Mapping[int, str], Sequence[int]]]] = None
    keep_initializers_as_inputs: Optional[bool] = None
    custom_opsets: Optional[Mapping[str, int]] = None
    export_modules_as_functions: Union[bool, Set[Type[torch.nn.Module]]] = False