INNER CODE UNIT · Python

_to_device

onnxsim/onnxsim · onnxsim/rpc/server.py:145

    def _to_device(self, inputs: Dict[str, np.ndarray]):
        return {
            k: self._Tensor(v, device=self.device).realize() for k, v in inputs.items()
        }

    def run(self, inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        outs = self._forward(self._to_device(inputs))
        return {name: out.numpy() for name, out in zip(self.output_names, outs)}

    def time(self, inputs: Dict[str, np.ndarray], number: int, repeat: int):
        """(per-call seconds, statistics), measuring steady-state ``TinyJit`` replays.

        Replaying a captured graph removes the Python ONNX interpreter's per-call cost, so the
        numbers reflect the generated kernels. ``stats`` also carries the eager (interpreted) call
        time and, from one replay under ``DEBUG=2``, device kernel time, FLOPs and bytes moved.
        """
        import contextlib
        import io

View source record →