torch.mtia.mtia_graph

Created On: Nov 10, 2025 | Last Updated On: Nov 10, 2025

The MTIA backend is implemented out of the tree, only interfaces are defined here.

class torch.mtia.mtia_graph.MTIAGraph(keep_graph=False)

Wrapper around a MTIA graph.

Return type:

Self

capture_begin(pool)

Begin capturing a MTIA graph.

capture_end()

End the capture of a MTIA graph.

instantiate()

Instantiate the captured MTIA graph.

pool()

Return an opaque token representing the id of this graph’s memory pool

Return type:

tuple[int, int]

replay()

Replay the captured MTIA graph.

reset()

Destroy the captured graph and reset the states.

class torch.mtia.mtia_graph.graph(mtia_graph, pool=None, stream=None)