TensorPlay
API reference
latest (dev)
Copy
View Markdown

Latest development documentation · Updated 2026-10-15

tensorplay.autograd.graph API

Functions 2

#

save_on_cpu

functionFull reference ↗
tensorplay.autograd.graph.save_on_cpu(pin_memory: bool = False, device_type: str = 'cuda')[source]

Store saved values in host memory and restore their original device.

#

saved_tensors_hooks

functionFull reference ↗
tensorplay.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

Context manager that installs a (pack, unpack) pair for every tensor saved via Function()’s save_for_backward while active.

  • pack_hook(tensor) runs once per saved tensor at save time and may return any object (e.g. a CPU copy, a key into external storage).

  • unpack_hook(packed) runs when ctx.saved_tensors reads it back and must return an equivalent tensorplay.Tensor.

If either hook raises, the exception is surfaced from the corresponding operation.

On this page

Ask DeepWiki